# Buggy code
def _scatter_add_transpose_rule(
    t,
    operand,
    scatter_indices,
    updates,
    update_jaxpr,
    update_consts,
    dimension_numbers,
    updates_shape,
):
    assert scatter_indices is not None
    operand_t = update_t = None
    if operand is None:
        operand_t = t

    if updates is None:
        gather_dnums = GatherDimensionNumbers(
            offset_dims=dimension_numbers.update_window_dims,
            collapsed_slice_dims=dimension_numbers.inserted_window_dims,
            start_index_map=dimension_numbers.scatter_dims_to_operand_dims,
        )
        slice_sizes = []
        pos = 0
        for i in xrange(len(t.shape)):
            if i in dimension_numbers.inserted_window_dims:
                slice_sizes.append(1)
            else:
                slice_sizes.append(
                    updates_shape[dimension_numbers.update_window_dims[pos]]
                )
                pos += 1
        update_t = gather(
            t, scatter_indices, dimension_numbers=gather_dnums, slice_sizes=slice_sizes
        )
    return [operand_t, None, update_t]


# Traceback (AttributeError)
Traceback (most recent call last):
File "/usr/local/google/home/schsam/.local/lib/python2.7/site-packages/absl/third_party/unittest3_backport/case.py", line 37, in testPartExecutor
yield
File "/usr/local/google/home/schsam/.local/lib/python2.7/site-packages/absl/third_party/unittest3_backport/case.py", line 162, in run
testMethod()
File "/usr/local/google/home/schsam/.local/lib/python2.7/site-packages/absl/testing/parameterized.py", line 262, in bound_param_test
test_method(self, **testcase_params)
File "simulate_test.py", line 164, in test_grad_through_nvt
assert grad(do_sim)(1.0) > 0.0
File "/usr/local/google/home/schsam/Source/jax/jax/api.py", line 235, in grad_f
_, g = value_and_grad_f(*args, **kwargs)
File "/usr/local/google/home/schsam/Source/jax/jax/api.py", line 289, in value_and_grad_f
g = vjp_py(onp.ones((), dtype=dtype))
File "/usr/local/google/home/schsam/Source/jax/jax/api_util.py", line 62, in apply_jaxtree_fun
ans = fun(*args)
File "/usr/local/google/home/schsam/Source/jax/jax/api.py", line 822, in out_vjp_packed
return out_vjp(cotangent_in)
File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 112, in vjp_
_, arg_cts = backward_pass(jaxpr, consts, (), dummy_args, dummy_primal_and_ct)
File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 180, in backward_pass
eqn.params, subjaxprs, sub_consts, sub_freevar_vals, invals, ct_in)
File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 536, in call_transpose
ans = primitive.bind(fun, all_args, **params)
File "/usr/local/google/home/schsam/Source/jax/jax/core.py", line 636, in call_bind
ans = primitive.impl(f, *args, **params)
File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/xla.py", line 591, in xla_call_impl
compiled_fun = xla_callable(fun, device_values, *map(abstractify, args))
File "/usr/local/google/home/schsam/Source/jax/jax/linear_util.py", line 208, in memoized_fun
ans = call(f, *args)
File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/xla.py", line 604, in xla_callable
jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
File "/usr/local/google/home/schsam/Source/jax/jax/linear_util.py", line 147, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/usr/local/google/home/schsam/Source/jax/jax/interpreters/ad.py", line 186, in backward_pass
cts_out = get_primitive_transpose(eqn.primitive)(ct_in, *invals, **eqn.params)
File "/usr/local/google/home/schsam/Source/jax/jax/lax/lax.py", line 2818, in _scatter_add_transpose_rule
for i in xrange(len(t.shape)):
AttributeError: 'Zero' object has no attribute 'shape'

# Fixed code
def _scatter_add_transpose_rule(
    t,
    operand,
    scatter_indices,
    updates,
    update_jaxpr,
    update_consts,
    dimension_numbers,
    updates_shape,
):
    assert scatter_indices is not None
    if t is ad_util.zero:
        return [ad_util.zero, None, ad_util.zero]

    operand_t = update_t = None
    if operand is None:
        operand_t = t

    if updates is None:
        gather_dnums = GatherDimensionNumbers(
            offset_dims=dimension_numbers.update_window_dims,
            collapsed_slice_dims=dimension_numbers.inserted_window_dims,
            start_index_map=dimension_numbers.scatter_dims_to_operand_dims,
        )
        slice_sizes = []
        pos = 0
        for i in xrange(len(t.shape)):
            if i in dimension_numbers.inserted_window_dims:
                slice_sizes.append(1)
            else:
                slice_sizes.append(
                    updates_shape[dimension_numbers.update_window_dims[pos]]
                )
                pos += 1
        update_t = gather(
            t, scatter_indices, dimension_numbers=gather_dnums, slice_sizes=slice_sizes
        )
    return [operand_t, None, update_t]