test_make_fx_no_decompose is failing on master

[samdow]

Currently covered with a skip, previously it was infinitely recursing on the prim decomposition

[samdow]

Full Stack Trace

Traceback (most recent call last):
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_device_type.py", line 389, in instantiated_test
    raise rte
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/testing/_internal/common_device_type.py", line 376, in instantiated_test
    result = test(self, **param_kwargs)
  File "test/test_pythonkey.py", line 122, in test_make_fx_no_decompose
    fx_f = make_fx(grad(f), decomposition_table)(torch.randn(5))
  File "/home/circleci/project/functorch/_src/python_key.py", line 206, in wrapped
    t = pythonkey_trace(wrap_key(f, args), concrete_args=tuple(phs))
  File "/home/circleci/project/functorch/_src/python_key.py", line 169, in pythonkey_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/_symbolic_trace.py", line 581, in trace
    self.create_node('output', 'output', (self.create_arg(fn(*args)),), {},
  File "/home/circleci/project/functorch/_src/python_key.py", line 188, in wrapped
    out = f(*tree_args)
  File "/home/circleci/project/functorch/_src/eager_transforms.py", line 1191, in wrapper
    results = grad_and_value(func, argnums, has_aux=has_aux)(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/eager_transforms.py", line 1083, in wrapper
    flat_grad_input = _autograd_grad(flat_outputs, flat_diff_args, create_graph=True)
  File "/home/circleci/project/functorch/_src/eager_transforms.py", line 102, in _autograd_grad
    grad_inputs = torch.autograd.grad(diff_outputs, inputs, grad_outputs,
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/autograd/__init__.py", line 276, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_decomp/decompositions.py", line 47, in inner
    r = f(*tree_map(increase_prec, args), **tree_map(increase_prec, kwargs))
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_decomp/decompositions.py", line 66, in tanh_backward
    return out_grad * (1 - y * y).conj_physical()
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 681, in _ref
    result = prim(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/functorch/_src/python_key.py", line 79, in __torch_dispatch__
    return CURRENT_DECOMPOSITION_TABLE[func_overload](*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 679, in _ref
    a, b = broadcast(a, b)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 428, in broadcast
    return tuple(map(lambda x: _maybe_broadcast(x, common_shape), args))
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 428, in <lambda>
    return tuple(map(lambda x: _maybe_broadcast(x, common_shape), args))
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_refs/__init__.py", line 422, in _maybe_broadcast
    return prims.broadcast_in_dim(x, common_shape, dims)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 174, in _prim
    return impl_aten(*args, **kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_prims/__init__.py", line 806, in _broadcast_in_dim_aten
    return v.expand(shape)
  File "/home/circleci/project/functorch/_src/python_key.py", line 91, in __torch_dispatch__
    proxy_out = func(*proxy_args, **proxy_kwargs)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/_ops.py", line 143, in __call__
    return self._op(*args, **kwargs or {})
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/proxy.py", line 315, in __torch_function__
    return tracer.create_proxy('call_function', orig_method, args, kwargs,
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/proxy.py", line 59, in create_proxy
    args_ = self.create_arg(args)
  File "/home/circleci/project/functorch/_src/python_key.py", line 162, in create_arg
    return super().create_arg(a)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/_symbolic_trace.py", line 291, in create_arg
    return super().create_arg(a)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/proxy.py", line 121, in create_arg
    return type(a)(self.create_arg(elem) for elem in a)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/proxy.py", line 121, in <genexpr>
    return type(a)(self.create_arg(elem) for elem in a)
  File "/home/circleci/project/functorch/_src/python_key.py", line 162, in create_arg
    return super().create_arg(a)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/_symbolic_trace.py", line 291, in create_arg
    return super().create_arg(a)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/proxy.py", line 121, in create_arg
    return type(a)(self.create_arg(elem) for elem in a)
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/fx/proxy.py", line 121, in <genexpr>
    return type(a)(self.create_arg(elem) for elem in a)
  File "/home/circleci/project/functorch/_src/python_key.py", line 146, in create_arg
    if isinstance(a, torch.nn.Parameter):
  File "/home/circleci/project/env/lib/python3.8/site-packages/torch/nn/parameter.py", line 10, in __instancecheck__
    return super().__instancecheck__(instance) or (
RecursionError: maximum recursion depth exceeded while calling a Python object