# Buggy code
def _conv_transpose_padding(k, s, padding):
    """Calculate before and after padding for a dim of transposed convolution.

    Args:
      k: int: kernel dimension.
      s: int: dimension stride value.
      padding: 'same' or 'valid' padding mode for original forward conv.

    Returns:
      2-tuple: ints: before and after padding for transposed convolution.
    """
    if padding == "SAME":
        pad_len = k + s - 2
        if s > k - 1:
            pad_a = k - 1
        else:
            pad_a = int(onp.ceil(pad_len / 2))
    elif padding == "VALID":
        pad_len = k + s - 2 + max(k - s, 0)
        pad_a = k - 1
    else:
        raise ValueError("Padding mode must be `SAME` or `VALID`.")
    pad_b = pad_len - pad_a
    return pad_a, pad_b


# 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 _conv_transpose_padding(k, s, padding):
    """Calculate before and after padding for a dim of transposed convolution.

    Args:
      k: int: kernel dimension.
      s: int: dimension stride value.
      padding: 'same' or 'valid' padding mode for original forward conv.

    Returns:
      2-tuple: ints: before and after padding for transposed convolution.
    """
    if padding == "SAME":
        pad_len = k + s - 2
        if s > k - 1:
            pad_a = k - 1
        else:
            pad_a = int(onp.ceil(pad_len / 2))
    elif padding == "VALID":
        pad_len = k + s - 2 + _max(k - s, 0)
        pad_a = k - 1
    else:
        raise ValueError("Padding mode must be `SAME` or `VALID`.")
    pad_b = pad_len - pad_a
    return pad_a, pad_b