jacobian接口问题：oneflow的_autograd_grad()函数对比pytorch的_autograd_grad()函数缺乏相关功能 about oneflow HOT 4 CLOSED

import torch as torch_original
import oneflow as flow

def exp_reducer(x):
    return x.exp().sum(dim=1)

print("torch result: ")
inputs = torch_original.rand(2,2,requires_grad=True)
outputs = exp_reducer(inputs)


torch_inputs = (inputs,)
torch_outputs = (outputs,)
torch_grad_outputs = (torch_original.eye(2),)
torch_result = torch_original.autograd.grad(
                    torch_outputs,
                    torch_inputs,
                    torch_grad_outputs,
                    allow_unused=True,
                    create_graph=False,
                    retain_graph=None,
                    is_grads_batched=True,
                )
print(torch_result)

print("oneflow result: ")
flow_inputs = (flow.tensor(inputs.detach().numpy(),requires_grad=True),)
flow_outputs = ( exp_reducer(flow_inputs[0]), )
flow_grad_outputs = (flow.eye(2),)
flow_result = flow.autograd.grad(
        flow_outputs,
        flow_inputs,
        flow_grad_outputs,
        allow_unused=True,
        create_graph=False,
        retain_graph=None,
    )
print(flow_result)

torch result: 
(tensor([[[1.0448, 2.2668],
         [0.0000, 0.0000]],

        [[0.0000, 0.0000],
         [2.5238, 1.5629]]]),)
oneflow result: 
Traceback (most recent call last):
  File "../../test/test.py", line 107, in <module>
    flow_result = flow.autograd.grad(
  File "/workspace/software/oneflow/python/oneflow/autograd/autograd.py", line 63, in grad
    in_grads = grad_api(
oneflow._oneflow_internal.exception.Exception: out_grad's shape must be same as output's ((2,) vs (2,2))
  File "oneflow/api/python/autograd/autograd.cpp", line 113, in Grad
    CheckAndInitOutGrads(outputs, out_grads)
  File "oneflow/api/python/autograd/autograd.cpp", line 73, in CheckAndInitOutGrads
    CHECK_OR_RETURN(*(outputs.at(i)->shape()) == *(out_grads.at(i)->shape()))
Error Type: oneflow.ErrorProto.check_failed_error