Convolution backward op missing

[samos123]

Observed when trying a training job that is doing convolution:

torchax/tensor.py", line 441, in _get_op_or_decomp
    raise OperatorNotFound(
torchax.tensor.OperatorNotFound: Operator with name aten::convolution_backward has no lowering

[qihqi]

How are you doing training? Please follow https://google.github.io/torchax/tutorials/trainingyt/ and use make_train_step to do training instead of traditional torch loss.backward() style training.

The difference is that the former will use jax.grad to compute gradients, and jax uses different formula for computing backward of conv that will not need the aten::convolution_backward op.