gurudesh/pytracebugs huggingface dataset:
# Buggy code
def eig_cpu_translation_rule(c, operand):
    out = lapack.jax_geev(c, operand)
    return c.Tuple(
        c.GetTupleElement(out, 0), c.GetTupleElement(out, 1), c.GetTupleElement(out, 2)
    )


# Traceback (numpy.linalg.LinAlgError)
/home/luke/.local/lib/python3.5/site-packages/jax/scipy/linalg.py:36: UserWarning: scipy.linalg support is experimental and may cause silent failures or wrong outputs
warnings.warn(_EXPERIMENTAL_WARNING)
[[0. 1. 2.]
[0. 4. 5.]
[0. 0. 8.]]
From Scipy
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
print(linalg.cholesky(x))
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 91, in cholesky
check_finite=check_finite)
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 40, in _cholesky
"definite" % info)
numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite

# Fixed code
def eig_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    batch_dims = shape.dimensions()[:-2]
    geev_out = lapack.jax_geev(c, operand)
    w = c.GetTupleElement(geev_out, 0)
    vl = c.GetTupleElement(geev_out, 1)
    vr = c.GetTupleElement(geev_out, 2)
    ok = c.Eq(c.GetTupleElement(geev_out, 3), c.ConstantS32Scalar(0))
    w = _broadcasting_select(
        c, c.Reshape(ok, None, batch_dims + (1,)), w, _nan_like(c, w)
    )
    vl = _broadcasting_select(
        c, c.Reshape(ok, None, batch_dims + (1, 1)), vl, _nan_like(c, vl)
    )
    vr = _broadcasting_select(
        c, c.Reshape(ok, None, batch_dims + (1, 1)), vr, _nan_like(c, vr)
    )
    return c.Tuple(w, vl, vr)
# Buggy code
def eigh_cpu_translation_rule(c, operand, lower):
    out = lapack.jax_syevd(c, operand, lower=lower)
    return c.Tuple(c.GetTupleElement(out, 0), c.GetTupleElement(out, 1))


# Traceback (numpy.linalg.LinAlgError)
/home/luke/.local/lib/python3.5/site-packages/jax/scipy/linalg.py:36: UserWarning: scipy.linalg support is experimental and may cause silent failures or wrong outputs
warnings.warn(_EXPERIMENTAL_WARNING)
[[0. 1. 2.]
[0. 4. 5.]
[0. 0. 8.]]
From Scipy
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
print(linalg.cholesky(x))
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 91, in cholesky
check_finite=check_finite)
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 40, in _cholesky
"definite" % info)
numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite

# Fixed code
def eigh_cpu_translation_rule(c, operand, lower):
    shape = c.GetShape(operand)
    batch_dims = shape.dimensions()[:-2]
    syevd_out = lapack.jax_syevd(c, operand, lower=lower)
    v = c.GetTupleElement(syevd_out, 0)
    w = c.GetTupleElement(syevd_out, 1)
    ok = c.Eq(c.GetTupleElement(syevd_out, 2), c.ConstantS32Scalar(0))
    v = _broadcasting_select(
        c, c.Reshape(ok, None, batch_dims + (1, 1)), v, _nan_like(c, v)
    )
    w = _broadcasting_select(
        c, c.Reshape(ok, None, batch_dims + (1,)), w, _nan_like(c, w)
    )
    return c.Tuple(v, w)
# Buggy code
def lu_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    out = lapack.jax_getrf(c, operand)
    lu = c.GetTupleElement(out, 0)
    # Subtract 1 from the pivot to get 0-based indices.
    pivot = c.Sub(c.GetTupleElement(out, 1), c.ConstantS32Scalar(1))
    # Throw away the `info` value, because we have no way to report errors.
    return c.Tuple(lu, pivot)


# Traceback (numpy.linalg.LinAlgError)
/home/luke/.local/lib/python3.5/site-packages/jax/scipy/linalg.py:36: UserWarning: scipy.linalg support is experimental and may cause silent failures or wrong outputs
warnings.warn(_EXPERIMENTAL_WARNING)
[[0. 1. 2.]
[0. 4. 5.]
[0. 0. 8.]]
From Scipy
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
print(linalg.cholesky(x))
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 91, in cholesky
check_finite=check_finite)
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 40, in _cholesky
"definite" % info)
numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite

# Fixed code
def lu_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    batch_dims = shape.dimensions()[:-2]
    getrf_out = lapack.jax_getrf(c, operand)
    lu = c.GetTupleElement(getrf_out, 0)
    # Subtract 1 from the pivot to get 0-based indices.
    pivot = c.Sub(c.GetTupleElement(getrf_out, 1), c.ConstantS32Scalar(1))
    ok = c.Eq(c.GetTupleElement(getrf_out, 2), c.ConstantS32Scalar(0))
    lu = _broadcasting_select(
        c, c.Reshape(ok, None, batch_dims + (1, 1)), lu, _nan_like(c, lu)
    )
    return c.Tuple(lu, pivot)
# Buggy code
def svd_cpu_translation_rule(c, operand, full_matrices, compute_uv):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
        out = lapack.jax_gesdd(
            c, operand, full_matrices=full_matrices, compute_uv=compute_uv
        )
        return c.Tuple(
            c.GetTupleElement(out, 0),
            c.GetTupleElement(out, 1),
            c.GetTupleElement(out, 2),
        )
    else:
        raise NotImplementedError(
            "Only unbatched singular value decomposition is implemented on CPU"
        )


# Traceback (numpy.linalg.LinAlgError)
/home/luke/.local/lib/python3.5/site-packages/jax/scipy/linalg.py:36: UserWarning: scipy.linalg support is experimental and may cause silent failures or wrong outputs
warnings.warn(_EXPERIMENTAL_WARNING)
[[0. 1. 2.]
[0. 4. 5.]
[0. 0. 8.]]
From Scipy
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
print(linalg.cholesky(x))
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 91, in cholesky
check_finite=check_finite)
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 40, in _cholesky
"definite" % info)
numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite

# Fixed code
def svd_cpu_translation_rule(c, operand, full_matrices, compute_uv):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
        gesdd_out = lapack.jax_gesdd(
            c, operand, full_matrices=full_matrices, compute_uv=compute_uv
        )
        s = c.GetTupleElement(gesdd_out, 0)
        u = c.GetTupleElement(gesdd_out, 1)
        vt = c.GetTupleElement(gesdd_out, 2)
        ok = c.Eq(c.GetTupleElement(gesdd_out, 3), c.ConstantS32Scalar(0))
        s = _broadcasting_select(c, c.Reshape(ok, None, (1,)), s, _nan_like(c, s))
        u = _broadcasting_select(c, c.Reshape(ok, None, (1, 1)), u, _nan_like(c, u))
        vt = _broadcasting_select(c, c.Reshape(ok, None, (1, 1)), vt, _nan_like(c, vt))
        return c.Tuple(s, u, vt)
    else:
        raise NotImplementedError(
            "Only unbatched singular value decomposition is implemented on CPU"
        )
# Buggy code
def cholesky_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
        potrf_output = lapack.jax_potrf(c, operand, lower=True)
        result = c.GetTupleElement(potrf_output, 0)
        info = c.GetTupleElement(potrf_output, 1)
        return c.Select(
            c.Eq(info, c.ConstantS32Scalar(0)),
            result,
            c.Broadcast(
                c.Constant(onp.array(onp.nan, dtype=dtype)), shape.dimensions()
            ),
        )
    else:
        # Fall back to the HLO implementation for batched Cholesky decomposition or
        # unsupported types.
        # TODO(phawkins): support LAPACK primitives in batched mode.
        return c.Cholesky(operand)


# Traceback (numpy.linalg.LinAlgError)
/home/luke/.local/lib/python3.5/site-packages/jax/scipy/linalg.py:36: UserWarning: scipy.linalg support is experimental and may cause silent failures or wrong outputs
warnings.warn(_EXPERIMENTAL_WARNING)
[[0. 1. 2.]
[0. 4. 5.]
[0. 0. 8.]]
From Scipy
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
print(linalg.cholesky(x))
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 91, in cholesky
check_finite=check_finite)
File "/home/luke/.local/lib/python3.5/site-packages/scipy/linalg/decomp_cholesky.py", line 40, in _cholesky
"definite" % info)
numpy.linalg.LinAlgError: 1-th leading minor of the array is not positive definite

# Fixed code
def cholesky_cpu_translation_rule(c, operand):
    shape = c.GetShape(operand)
    dtype = shape.element_type().type
    if len(shape.dimensions()) == 2 and dtype in _cpu_lapack_types:
        potrf_output = lapack.jax_potrf(c, operand, lower=True)
        result = c.GetTupleElement(potrf_output, 0)
        info = c.GetTupleElement(potrf_output, 1)
        return c.Select(
            c.Eq(info, c.ConstantS32Scalar(0)), result, _nan_like(c, result)
        )
    else:
        # Fall back to the HLO implementation for batched Cholesky decomposition or
        # unsupported types.
        # TODO(phawkins): support LAPACK primitives in batched mode.
        return c.Cholesky(operand)
# Buggy code
def _wrap_hashably(arg):
    try:
        hash(arg)
    except TypeError:
        return WrapHashably(arg)
    else:
        return Hashable(arg)


# Traceback (ValueError)
Traceback (most recent call last):

File "<ipython-input-7-563ff9ef5fe4>", line 1, in <module>
runfile('/home/hpc/capm/sn0523/SIR/minExample.py', wdir='/home/hpc/capm/sn0523/SIR')

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/spyder_kernels/customize/spydercustomize.py", line 827, in runfile
execfile(filename, namespace)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/spyder_kernels/customize/spydercustomize.py", line 110, in execfile
exec(compile(f.read(), filename, 'exec'), namespace)

File "/home/hpc/capm/sn0523/SIR/minExample.py", line 25, in <module>
main()

File "/home/hpc/capm/sn0523/SIR/minExample.py", line 21, in main
secondCall = f(x,v)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/api.py", line 123, in f_jitted
out = xla.xla_call(flat_fun, *args_flat, device_values=device_values)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/core.py", line 658, in call_bind
ans = primitive.impl(f, *args, **params)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/interpreters/xla.py", line 653, in xla_call_impl
compiled_fun = xla_callable(fun, device_values, *map(abstractify, args))

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/linear_util.py", line 201, in memoized_fun
if key in cache:

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/linear_util.py", line 174, in __eq__
return self.hashable_payload() == other.hashable_payload()

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/interpreters/xla.py", line 489, in forward_method
return fun(getattr(self, attrname), *args)

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

# Fixed code
def _wrap_hashably(arg):
    try:
        hash(arg)
    except TypeError:
        return WrapHashably(arg)  # e.g. ndarrays, DeviceArrays
    else:
        return Hashable(arg)
# Buggy code
def __hash__(self):
    # TODO(mattjj): this is not semantically correct because it is possible
    # __eq__ is true for values with unequal __hash__ values. However, the
    # main use case at the moment is memoization for which false negatives are
    # fine.
    return id(self)


# Traceback (ValueError)
Traceback (most recent call last):

File "<ipython-input-7-563ff9ef5fe4>", line 1, in <module>
runfile('/home/hpc/capm/sn0523/SIR/minExample.py', wdir='/home/hpc/capm/sn0523/SIR')

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/spyder_kernels/customize/spydercustomize.py", line 827, in runfile
execfile(filename, namespace)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/spyder_kernels/customize/spydercustomize.py", line 110, in execfile
exec(compile(f.read(), filename, 'exec'), namespace)

File "/home/hpc/capm/sn0523/SIR/minExample.py", line 25, in <module>
main()

File "/home/hpc/capm/sn0523/SIR/minExample.py", line 21, in main
secondCall = f(x,v)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/api.py", line 123, in f_jitted
out = xla.xla_call(flat_fun, *args_flat, device_values=device_values)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/core.py", line 658, in call_bind
ans = primitive.impl(f, *args, **params)

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/interpreters/xla.py", line 653, in xla_call_impl
compiled_fun = xla_callable(fun, device_values, *map(abstractify, args))

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/linear_util.py", line 201, in memoized_fun
if key in cache:

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/linear_util.py", line 174, in __eq__
return self.hashable_payload() == other.hashable_payload()

File "/home/hpc/capm/sn0523/jax/lib/python3.6/site-packages/jax/interpreters/xla.py", line 489, in forward_method
return fun(getattr(self, attrname), *args)

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

# Fixed code
def __hash__(self):
    raise TypeError("JAX DeviceArray, like numpy.ndarray, is not hashable.")
# Buggy code
def arange(*args, **kwargs):
    dtype = kwargs.pop("dtype", None)
    if not args:
        raise TypeError(
            "Required argument 'start' (pos 1) not found"
        )  # same as numpy error

    # If called like np.arange(N), we create a lazy lax._IotaConstant.
    if len(args) == 1 and not kwargs:
        (stop,) = args
        dtype = dtype or _dtype(stop)
        if onp.issubdtype(dtype, onp.integer):
            return lax.iota(dtype, stop)  # avoids materializing

    # Fall back to instantiating an ndarray in host memory
    return onp.arange(*args, **kwargs)


# Traceback (AssertionError)
int64
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-12-5b101f76f4c3> in <module>()
1 a = np.arange(4, dtype=np.complex64)
2 print(a.dtype)
----> 3 assert a.dtype is np.complex64

AssertionError:

# Fixed code
def arange(*args, **kwargs):
    dtype = kwargs.get("dtype", None)
    if not args:
        raise TypeError(
            "Required argument 'start' (pos 1) not found"
        )  # same as numpy error

    # If called like np.arange(N), we create a lazy lax._IotaConstant.
    if len(args) == 1 and not kwargs:
        (stop,) = args
        dtype = dtype or _dtype(stop)
        if onp.issubdtype(dtype, onp.integer):
            return lax.iota(dtype, stop)  # avoids materializing

    # Fall back to instantiating an ndarray in host memory
    return onp.arange(*args, **kwargs)
# Buggy code
def partial_eval(f, trace, pvs):
    f = trace_to_subjaxpr(f, trace.master)
    return partial_eval_wrapper(f, tuple(pvs))


# Traceback (RuntimeError)
$ python loop.py
/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py:144: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
File "loop.py", line 11, in <module>
print(lax.while_loop(cond, body, (33, 4)))
File "/usr/local/google/home/phawkins/p/jax/jax/lax/lax_control_flow.py", line 108, in while_loop
aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
File "/usr/local/google/home/phawkins/p/jax/jax/core.py", line 75, in bind
return self.impl(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 51, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 57, in xla_primitive_callable
built_c = primitive_computation(prim, *shapes, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 74, in primitive_computation
raise e
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 70, in primitive_computation
return c.Build()
File "/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py", line 283, in Build
*args, **kwargs)
File "/usr/local/google/home/phawkins/.pyenv/versions/3.7.3/lib/python3.7/site-packages/jaxlib/xla_client.py", line 875, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> pred[]; body: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> (((), s32[]), (s32[]), (s32[])); init: ((s32[], s32[]), (s32[]), (s32[]))..:

# Fixed code
def partial_eval(f, trace, pvs):
    f = trace_to_subjaxpr(f, trace.master, False)
    return partial_eval_wrapper(f, tuple(pvs))
# Buggy code
def trace_to_jaxpr(fun, pvals):
    """Traces a function, given abstract inputs, to a jaxpr."""
    with new_master(JaxprTrace) as master:
        fun = trace_to_subjaxpr(fun, master)
        jaxpr, (out_pval, consts, env) = fun.call_wrapped(pvals)
        assert not env
        del master

    return jaxpr, out_pval, consts


# Traceback (RuntimeError)
$ python loop.py
/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py:144: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
File "loop.py", line 11, in <module>
print(lax.while_loop(cond, body, (33, 4)))
File "/usr/local/google/home/phawkins/p/jax/jax/lax/lax_control_flow.py", line 108, in while_loop
aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
File "/usr/local/google/home/phawkins/p/jax/jax/core.py", line 75, in bind
return self.impl(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 51, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 57, in xla_primitive_callable
built_c = primitive_computation(prim, *shapes, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 74, in primitive_computation
raise e
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 70, in primitive_computation
return c.Build()
File "/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py", line 283, in Build
*args, **kwargs)
File "/usr/local/google/home/phawkins/.pyenv/versions/3.7.3/lib/python3.7/site-packages/jaxlib/xla_client.py", line 875, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> pred[]; body: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> (((), s32[]), (s32[]), (s32[])); init: ((s32[], s32[]), (s32[]), (s32[]))..:

# Fixed code
def trace_to_jaxpr(fun, pvals, **kwargs):
    """Traces a function, given abstract inputs, to a jaxpr."""
    instantiate = kwargs.pop("instantiate", False)
    with new_master(JaxprTrace) as master:
        fun = trace_to_subjaxpr(fun, master, instantiate)
        jaxpr, (out_pval, consts, env) = fun.call_wrapped(pvals)
        assert not env
        del master

    return jaxpr, out_pval, consts
# Buggy code
def trace_to_subjaxpr(master, pvals):
    assert all([isinstance(pv, PartialVal) for pv in pvals]), pvals
    trace = JaxprTrace(master, core.cur_sublevel())
    in_tracers = map(trace.new_arg, pvals)
    out_tracer = yield in_tracers, {}
    out_tracer = trace.full_raise(out_tracer)
    jaxpr, consts, env = tracers_to_jaxpr(in_tracers, out_tracer)
    out_pval = out_tracer.pval
    del trace, in_tracers, out_tracer
    yield jaxpr, (out_pval, consts, env)


# Traceback (RuntimeError)
$ python loop.py
/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py:144: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
File "loop.py", line 11, in <module>
print(lax.while_loop(cond, body, (33, 4)))
File "/usr/local/google/home/phawkins/p/jax/jax/lax/lax_control_flow.py", line 108, in while_loop
aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
File "/usr/local/google/home/phawkins/p/jax/jax/core.py", line 75, in bind
return self.impl(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 51, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 57, in xla_primitive_callable
built_c = primitive_computation(prim, *shapes, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 74, in primitive_computation
raise e
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 70, in primitive_computation
return c.Build()
File "/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py", line 283, in Build
*args, **kwargs)
File "/usr/local/google/home/phawkins/.pyenv/versions/3.7.3/lib/python3.7/site-packages/jaxlib/xla_client.py", line 875, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> pred[]; body: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> (((), s32[]), (s32[]), (s32[])); init: ((s32[], s32[]), (s32[]), (s32[]))..:

# Fixed code
def trace_to_subjaxpr(master, instantiate, pvals):
    assert all([isinstance(pv, PartialVal) for pv in pvals]), pvals
    trace = JaxprTrace(master, core.cur_sublevel())
    in_tracers = map(trace.new_arg, pvals)
    out_tracer = yield in_tracers, {}
    out_tracer = trace.full_raise(out_tracer)

    if instantiate:
        out_tracer = trace.instantiate_const(out_tracer)

    jaxpr, consts, env = tracers_to_jaxpr(in_tracers, out_tracer)
    out_pval = out_tracer.pval
    del trace, in_tracers, out_tracer
    yield jaxpr, (out_pval, consts, env)
# Buggy code
def compiled_call_impl(fun, *args):
    with new_master(JaxprTrace, True) as master:
        pvals = map(abstractify, args)
        jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master).call_wrapped(pvals)
        jaxpr_ans = eval_jaxpr_raw(jaxpr, consts, env, *args)
        ans = merge_pvals(jaxpr_ans, pval)
        del master, pvals, pval, consts, env, jaxpr_ans, jaxpr
        return ans


# Traceback (RuntimeError)
$ python loop.py
/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py:144: UserWarning: No GPU/TPU found, falling back to CPU.
warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
File "loop.py", line 11, in <module>
print(lax.while_loop(cond, body, (33, 4)))
File "/usr/local/google/home/phawkins/p/jax/jax/lax/lax_control_flow.py", line 108, in while_loop
aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
File "/usr/local/google/home/phawkins/p/jax/jax/core.py", line 75, in bind
return self.impl(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 51, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 57, in xla_primitive_callable
built_c = primitive_computation(prim, *shapes, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 74, in primitive_computation
raise e
File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 70, in primitive_computation
return c.Build()
File "/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py", line 283, in Build
*args, **kwargs)
File "/usr/local/google/home/phawkins/.pyenv/versions/3.7.3/lib/python3.7/site-packages/jaxlib/xla_client.py", line 875, in Build
return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> pred[]; body: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> (((), s32[]), (s32[]), (s32[])); init: ((s32[], s32[]), (s32[]), (s32[]))..:

# Fixed code
def compiled_call_impl(fun, *args):
    with new_master(JaxprTrace, True) as master:
        pvals = map(abstractify, args)
        jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master, False).call_wrapped(
            pvals
        )
        jaxpr_ans = eval_jaxpr_raw(jaxpr, consts, env, *args)
        ans = merge_pvals(jaxpr_ans, pval)
        del master, pvals, pval, consts, env, jaxpr_ans, jaxpr
        return ans
# Buggy code
def atleast_1d(*arys):
    if len(arys) == 1:
        arr = array(arys[0])
        return arr if arr.ndim >= 1 else arr.reshape(-1)
    else:
        return [atleast_1d(arr) for arr in arys]


# Traceback (AttributeError)
In [3]: onp.atleast_1d(1)
Out[3]: array([1])

In [4]: np.atleast_1d(1)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-84084c6642da> in <module>
----> 1 np.atleast_1d(1)

~/src/jax/jax/numpy/lax_numpy.py in atleast_1d(*arys)
1182     arr = array(arys[0])
-> 1183     return arr if arr.ndim >= 1 else arr.reshape(-1)
1184   else:
1185     return [atleast_1d(arr) for arr in arys]

AttributeError: 'int' object has no attribute 'ndim'

In [5]: np.atleast_1d(np.int64(1))
Out[5]: array([1])

# Fixed code
def atleast_1d(*arys):
    if len(arys) == 1:
        arr = array(arys[0])
        return arr if ndim(arr) >= 1 else reshape(arr, -1)
    else:
        return [atleast_1d(arr) for arr in arys]
# Buggy code
def atleast_2d(*arys):
    if len(arys) == 1:
        arr = array(arys[0])
        return arr if arr.ndim >= 2 else arr.reshape((1, -1))
    else:
        return [atleast_2d(arr) for arr in arys]


# Traceback (AttributeError)
In [3]: onp.atleast_1d(1)
Out[3]: array([1])

In [4]: np.atleast_1d(1)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-84084c6642da> in <module>
----> 1 np.atleast_1d(1)

~/src/jax/jax/numpy/lax_numpy.py in atleast_1d(*arys)
1182     arr = array(arys[0])
-> 1183     return arr if arr.ndim >= 1 else arr.reshape(-1)
1184   else:
1185     return [atleast_1d(arr) for arr in arys]

AttributeError: 'int' object has no attribute 'ndim'

In [5]: np.atleast_1d(np.int64(1))
Out[5]: array([1])

# Fixed code
def atleast_2d(*arys):
    if len(arys) == 1:
        arr = array(arys[0])
        return arr if ndim(arr) >= 2 else reshape(arr, (1, -1))
    else:
        return [atleast_2d(arr) for arr in arys]
# Buggy code
def atleast_3d(*arys):
    if len(arys) == 1:
        arr = array(arys[0])
        if ndim(arr) <= 1:
            arr = arr.reshape((1, -1, 1))
        elif ndim(arr) == 2:
            arr = arr.reshape(shape(arr) + (1,))
        return arr
    else:
        return [atleast_3d(arr) for arr in arys]


# Traceback (AttributeError)
In [3]: onp.atleast_1d(1)
Out[3]: array([1])

In [4]: np.atleast_1d(1)
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-84084c6642da> in <module>
----> 1 np.atleast_1d(1)

~/src/jax/jax/numpy/lax_numpy.py in atleast_1d(*arys)
1182     arr = array(arys[0])
-> 1183     return arr if arr.ndim >= 1 else arr.reshape(-1)
1184   else:
1185     return [atleast_1d(arr) for arr in arys]

AttributeError: 'int' object has no attribute 'ndim'

In [5]: np.atleast_1d(np.int64(1))
Out[5]: array([1])

# Fixed code
def atleast_3d(*arys):
    if len(arys) == 1:
        arr = array(arys[0])
        if ndim(arr) <= 1:
            arr = reshape(arr, (1, -1, 1))
        elif ndim(arr) == 2:
            arr = reshape(arr, shape(arr) + (1,))
        return arr
    else:
        return [atleast_3d(arr) for arr in arys]
# Buggy code
def _normalize_by_window_size(dims, strides, padding):
    def rescale(outputs, inputs):
        one = np.ones(inputs.shape[1:3], dtype=inputs.dtype)
        window_sizes = lax.reduce_window(one, 0.0, lax.add, dims, strides, padding)
        return outputs / window_sizes

    return rescale


# Traceback (ValueError)
(-1, 16, 16, 3)
()
Traceback (most recent call last):
File "minimal_example.py", line 9, in <module>
apply_fun(params, np.zeros((100, 32, 32, 3)))
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/experimental/stax.py", line 172, in apply_fun
return rescale(out, inputs) if rescale else out
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/experimental/stax.py", line 183, in rescale
return outputs / window_sizes
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 350, in true_divide
x1, x2 = _promote_shapes(x1, x2)
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 134, in _promote_shapes
nd = len(_broadcast_shapes(*shapes))
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/util.py", line 161, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 151, in _broadcast_shapes
.format(tuple(map(tuple, shapes))))
ValueError: Incompatible shapes for broadcasting: ((100, 16, 16, 3), (1, 1, 16, 16))

# Fixed code
def _normalize_by_window_size(dims, strides, padding):
    def rescale(outputs, inputs):
        one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
        window_sizes = lax.reduce_window(one, 0.0, lax.add, dims, strides, padding)
        return outputs / window_sizes[..., np.newaxis]

    return rescale
# Buggy code
def rescale(outputs, inputs):
    one = np.ones(inputs.shape[1:3], dtype=inputs.dtype)
    window_sizes = lax.reduce_window(one, 0.0, lax.add, dims, strides, padding)
    return outputs / window_sizes


# Traceback (ValueError)
(-1, 16, 16, 3)
()
Traceback (most recent call last):
File "minimal_example.py", line 9, in <module>
apply_fun(params, np.zeros((100, 32, 32, 3)))
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/experimental/stax.py", line 172, in apply_fun
return rescale(out, inputs) if rescale else out
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/experimental/stax.py", line 183, in rescale
return outputs / window_sizes
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 350, in true_divide
x1, x2 = _promote_shapes(x1, x2)
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 134, in _promote_shapes
nd = len(_broadcast_shapes(*shapes))
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/util.py", line 161, in memoized_fun
ans = cache[key] = fun(*args, **kwargs)
File "/gpfs01/bethge/home/jrauber/PhD/063_jax/jax/jax/numpy/lax_numpy.py", line 151, in _broadcast_shapes
.format(tuple(map(tuple, shapes))))
ValueError: Incompatible shapes for broadcasting: ((100, 16, 16, 3), (1, 1, 16, 16))

# Fixed code
def rescale(outputs, inputs):
    one = np.ones(inputs.shape[1:-1], dtype=inputs.dtype)
    window_sizes = lax.reduce_window(one, 0.0, lax.add, dims, strides, padding)
    return outputs / window_sizes[..., np.newaxis]
# Buggy code
def bitcast_convert_type(operand, new_dtype):
    return bitcast_convert_type_p.bind(operand, new_dtype=new_dtype)


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def bitcast_convert_type(operand, new_dtype):
    new_dtype = xla_bridge.canonicalize_dtype(new_dtype)
    old_dtype = _dtype(operand)
    if old_dtype != new_dtype:
        return bitcast_convert_type_p.bind(operand, new_dtype=new_dtype)
    else:
        return operand
# Buggy code
def reduce(operand, init_value, computation, dimensions):
    monoid_reducer = _get_monoid_reducer(computation, init_value)
    if monoid_reducer:
        return monoid_reducer(operand, dimensions)
    else:
        jaxpr, consts = _reduction_jaxpr(computation, init_value)
        return reduce_p.bind(
            operand,
            init_value,
            jaxpr=jaxpr,
            consts=consts,
            dimensions=tuple(dimensions),
        )


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def reduce(operand, init_value, computation, dimensions):
    monoid_reducer = _get_monoid_reducer(computation, init_value)
    if monoid_reducer:
        return monoid_reducer(operand, dimensions)
    else:
        jaxpr, consts = _reduction_jaxpr(computation, init_value)
        return reduce_p.bind(
            operand,
            init_value,
            computation=computation,
            jaxpr=jaxpr,
            consts=consts,
            dimensions=tuple(dimensions),
        )
# Buggy code
def _get_monoid_reducer(monoid_op, x):
    aval = core.get_aval(x)
    if (type(aval) is ConcreteArray) and aval.shape == ():
        if monoid_op is add:
            return aval.val == 0 and _reduce_sum
        elif monoid_op is max:
            return aval.val == _get_max_identity(aval.dtype) and _reduce_max
        elif monoid_op is min:
            return aval.val == _get_min_identity(aval.dtype) and _reduce_min


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def _get_monoid_reducer(monoid_op, x):
    aval = core.get_aval(x)
    if (type(aval) is ConcreteArray) and aval.shape == ():
        if monoid_op is add:
            return aval.val == 0 and _reduce_sum
        elif monoid_op is max:
            return aval.val == _get_max_identity(aval.dtype) and _reduce_max
        elif monoid_op is min:
            return aval.val == _get_min_identity(aval.dtype) and _reduce_min
        elif monoid_op is bitwise_or and aval.dtype == onp.bool_:
            return aval.val == _get_max_identity(aval.dtype) and _reduce_or
        elif monoid_op is bitwise_and and aval.dtype == onp.bool_:
            return aval.val == _get_min_identity(aval.dtype) and _reduce_and
# Buggy code
def _get_max_identity(dtype):
    if onp.issubdtype(dtype, onp.floating):
        return onp.array(-onp.inf, dtype)
    elif onp.issubdtype(dtype, onp.integer):
        return onp.array(onp.iinfo(dtype).min, dtype)


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def _get_max_identity(dtype):
    if onp.issubdtype(dtype, onp.floating):
        return onp.array(-onp.inf, dtype)
    elif onp.issubdtype(dtype, onp.integer):
        return onp.array(onp.iinfo(dtype).min, dtype)
    elif onp.issubdtype(dtype, onp.bool_):
        return onp.array(False, onp.bool_)
# Buggy code
def _get_min_identity(dtype):
    if onp.issubdtype(dtype, onp.floating):
        return onp.array(onp.inf, dtype)
    elif onp.issubdtype(dtype, onp.integer):
        return onp.array(onp.iinfo(dtype).max, dtype)


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def _get_min_identity(dtype):
    if onp.issubdtype(dtype, onp.floating):
        return onp.array(onp.inf, dtype)
    elif onp.issubdtype(dtype, onp.integer):
        return onp.array(onp.iinfo(dtype).max, dtype)
    elif onp.issubdtype(dtype, onp.bool_):
        return onp.array(True, onp.bool_)
# Buggy code
def pad_batch_rule(batched_args, batch_dims, padding_config):
    operand, padding_value = batched_args
    operand_bdim, padding_value_bdim = batch_dims
    if padding_value_bdim is None:
        assert operand_bdim is not None
        padding_config = list(padding_config)
        padding_config.insert(operand_bdim, (0, 0, 0))
        return pad(operand, padding_value, padding_config), operand_bdim
    else:
        raise NotImplementedError


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def pad_batch_rule(batched_args, batch_dims, padding_config):
    operand, padding_value = batched_args
    operand_bdim, padding_value_bdim = batch_dims
    if padding_value_bdim is None:
        assert operand_bdim is not None
        padding_config = list(padding_config)
        padding_config.insert(operand_bdim, (0, 0, 0))
        return pad(operand, padding_value, padding_config), operand_bdim
    else:
        raise NotImplementedError  # loop and stack
# Buggy code
def reduce_shape_rule(operand, init_value, jaxpr, consts, dimensions):
    return tuple(onp.delete(operand.shape, dimensions))


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def reduce_shape_rule(operand, init_value, computation, jaxpr, consts, dimensions):
    return tuple(onp.delete(operand.shape, dimensions))
# Buggy code
def reduce_translation_rule(c, operand, init_value, jaxpr, consts, dimensions):
    xla_computation = _reduction_computation(c, jaxpr, consts, init_value)
    return c.Reduce(operand, init_value, xla_computation, dimensions)


# Traceback (TypeError)
import jax.numpy as np
from jax import vmap
vmap(np.any)(np.array([[True, False], [False, False]]))
jax/lib/xla_bridge.py:138: UserWarning: No GPU found, falling back to CPU.
warnings.warn('No GPU found, falling back to CPU.')
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "jax/api.py", line 149, in batched_fun
out_flat = batching.batch(flat_fun, in_flat, in_axes_, out_axes)
File "jax/interpreters/batching.py", line 43, in batch
out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
File "jax/linear_util.py", line 85, in call_wrapped
ans = self.f(*args, **self.kwargs)
File "jax/numpy/lax_numpy.py", line 607, in reduction
result = lax.reduce(a, _reduction_init_val(a, init_val), op, dims)
File "jax/lax.py", line 260, in reduce
dimensions=tuple(dimensions))
File "jax/core.py", line 74, in bind
out_tracer = top_trace.process_primitive(self, tracers, kwargs)
File "jax/interpreters/batching.py", line 119, in process_primitive
val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
TypeError: reducer_batcher() takes exactly 4 arguments (3 given)

# Fixed code
def reduce_translation_rule(
    c, operand, init_value, computation, jaxpr, consts, dimensions
):
    xla_computation = _reduction_computation(c, jaxpr, consts, init_value)
    return c.Reduce(operand, init_value, xla_computation, dimensions)