感谢阅读我的问题！
我刚刚开始学习Jax中的自定义梯度函数，我发现Jax在定义自定义函数时的方法非常优雅。
但有一件事让我感到困扰。
我创建了一个包装器，使lax卷积看起来像PyTorch的conv2d。
from jax import numpy as jnpfrom jax.random import PRNGKey, normal from jax import laxfrom torch.nn.modules.utils import _ntupleimport jaxfrom jax.nn.initializers import normalfrom jax import gradtorch_dims = {0: ('NC', 'OI', 'NC'), 1: ('NCH', 'OIH', 'NCH'), 2: ('NCHW', 'OIHW', 'NCHW'), 3: ('NCHWD', 'OIHWD', 'NCHWD')}def conv(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): n = len(input.shape) - 2 if type(stride) == int: stride = _ntuple(n)(stride) if type(padding) == int: padding = [(i, i) for i in _ntuple(n)(padding)] if type(dilation) == int: dilation = _ntuple(n)(dilation) return lax.conv_general_dilated(lhs=input, rhs=weight, window_strides=stride, padding=padding, lhs_dilation=dilation, rhs_dilation=None, dimension_numbers=torch_dims[n], feature_group_count=1, batch_group_count=1, precision=None, preferred_element_type=None)问题是我找不到使用其梯度函数的方法：
init = normal()rng = PRNGKey(42)x = init(rng, [128, 3, 224, 224])k = init(rng, [64, 3, 3, 3])y = conv(x, k)grad(conv)(y, k)这是我得到的结果：
ValueError: conv_general_dilated lhs feature dimension size divided by feature_group_count must equal the rhs input feature dimension size, but 64 // 1 != 3.请帮助我！
回答：
当我使用最新的jax和jaxlib版本（jax==0.2.22；jaxlib==0.1.72）运行你的代码时，我看到了以下错误：
TypeError: Gradient only defined for scalar-output functions. Output had shape: (128, 64, 222, 222).如果我创建一个使用conv的标量输出函数，梯度似乎可以正常工作：
result = grad(lambda x, k: conv(x, k).sum())(x, k)print(result.shape)# (128, 3, 224, 224)如果你使用的是较旧版本的JAX，你可以尝试更新到较新的版本——也许你看到的错误是由于已修复的bug引起的。
