Issue 1:
LayerError Traceback (most recent call last)

in

1 # this takes around 1 minute

2 # print(next(data_generator(Q1_test,Q2_test,512,pad=vocab[''],shuffle=False)))

----> 3 accuracy = classify(Q1_test,Q2_test, y_test, 0.7, model, vocab, batch_size = 512)

4 print("Accuracy", accuracy)

in classify(test_Q1, test_Q2, y, threshold, model, vocab, data_generator, batch_size)

27 # print(f'shape of y_test = {y_test.shape}\n')

28 # Call the model

---> 29 v1, v2 = model((q1,q2))

30

31 for j in range(batch_size):

~/anaconda3/lib/python3.8/site-packages/trax/layers/base.py in call(self, x, weights, state, rng)

171 self.state = state # Needed if the model wasn't fully initialized.

172 state = self.state

--> 173 outputs, new_state = self.pure_fn(x, weights, state, rng)

174 self.state = new_state

175 self.weights = weights

~/anaconda3/lib/python3.8/site-packages/trax/layers/base.py in pure_fn(self, x, weights, state, rng, use_cache)

520 # Skipping 3 lines as it's always the uninteresting internal call.

521 name, trace = self._name, _short_traceback(skip=3)

--> 522 raise LayerError(name, 'pure_fn',

523 self._caller, signature(x), trace) from None

524

LayerError: Exception passing through layer Parallel (in pure_fn):

layer created in file [...]/, line 27

layer input shapes: (ShapeDtype{shape:(512, 64), dtype:int64}, ShapeDtype{shape:(512, 64), dtype:int64})

File [...]/trax/layers/combinators.py, line 209, in forward

sub_outputs, sub_state = layer.pure_fn(x, w, s, r, use_cache=True)

LayerError: Exception passing through layer Serial (in pure_fn):

layer created in file [...]/, line 17

layer input shapes: ShapeDtype{shape:(512, 64), dtype:int64}

File [...]/trax/layers/combinators.py, line 88, in forward

outputs, s = layer.pure_fn(inputs, w, s, rng, use_cache=True)

LayerError: Exception passing through layer Embedding_41699_128 (in pure_fn):

layer created in file [...]/, line 18

layer input shapes: ShapeDtype{shape:(512, 64), dtype:int64}

File [...]/trax/layers/core.py, line 150, in forward

return jnp.take(self.weights, x, axis=0)

File [...]/jax/numpy/lax_numpy.py, line 3422, in take

return lax.gather(a, indices[..., None], dimension_numbers=dnums,

File [...]/jax/lax/lax.py, line 807, in gather

return gather_p.bind(

File [...]/site-packages/jax/core.py, line 276, in bind

return self.impl(*args, **kwargs)

File [...]/jax/interpreters/xla.py, line 224, in apply_primitive

compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)

File [...]/jax/interpreters/xla.py, line 257, in xla_primitive_callable

built_c = primitive_computation(prim, AxisEnv(nreps, (), (), None), backend,

File [...]/jax/interpreters/xla.py, line 316, in primitive_computation

raise RuntimeError(msg) from e

RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.:

This is a bug in JAX's shape-checking rules; please report it!

Issue 2:
code:

def gumbel_sample(log_probs, temperature=1.0):
"""Gumbel sampling from a categorical distribution."""
u = numpy.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
g = -np.log(-np.log(u))
return np.argmax(log_probs + g * temperature, axis=-1)

def predict(num_chars, prefix):
inp = [ord(c) for c in prefix]
result = [c for c in prefix]
max_len = len(prefix) + num_chars
for _ in range(num_chars):
cur_inp = np.array(inp + [0] * (max_len - len(inp)))
outp = model(cur_inp[None, :]) # Add batch dim.
next_char = gumbel_sample(outp[0, len(inp)])
inp += [int(next_char)]

if inp[-1] == 1:
break # EOS
result.append(chr(int(next_char)))

return "".join(result)

print(predict(32, ""))

Error:

LayerError Traceback (most recent call last)
<ipython-input-27-6f9f9afc30e6> in <module>
22 return "".join(result)
23
---> 24 print(predict(32, ""))

<ipython-input-27-6f9f9afc30e6> in predict(num_chars, prefix)
12 for _ in range(num_chars):
13 cur_inp = np.array(inp + [0] * (max_len - len(inp)))
---> 14 outp = model(cur_inp[None, :]) # Add batch dim.
15 next_char = gumbel_sample(outp[0, len(inp)])
16 inp += [int(next_char)]

/opt/conda/lib/python3.7/site-packages/trax/layers/base.py in __call__(self, x, weights, state, rng)
165 self.state = state # Needed if the model wasn't fully initialized.
166 state = self.state
--> 167 outputs, new_state = self.pure_fn(x, weights, state, rng)
168 self.state = new_state
169 self.weights = weights

/opt/conda/lib/python3.7/site-packages/trax/layers/base.py in pure_fn(self, x, weights, state, rng, use_cache)
448 name, trace = self._name, _short_traceback(skip=3)
449 raise LayerError(name, 'pure_fn',
--> 450 self._caller, signature(x), trace) from None
451
452 def output_signature(self, input_signature):

LayerError: Exception passing through layer Serial (in pure_fn):
layer created in file [...]/<ipython-input-13-6f7ffe0c061e>, line 21
layer input shapes: ShapeDtype{shape:(1, 32), dtype:int32}

File [...]/trax/layers/combinators.py, line 88, in forward
outputs, s = layer.pure_fn(inputs, w, s, rng, use_cache=True)

LayerError: Exception passing through layer Embedding_256_512 (in pure_fn):
layer created in file [...]/<ipython-input-13-6f7ffe0c061e>, line 18
layer input shapes: ShapeDtype{shape:(1, 32), dtype:int32}

File [...]/trax/layers/core.py, line 150, in forward
return jnp.take(self.weights, x, axis=0)

File [...]/jax/numpy/lax_numpy.py, line 3298, in take
slice_sizes=tuple(slice_sizes))

File [...]/jax/lax/lax.py, line 835, in gather
slice_sizes=canonicalize_shape(slice_sizes))

File [...]/site-packages/jax/core.py, line 273, in bind
return self.impl(*args, **kwargs)

File [...]/jax/interpreters/xla.py, line 228, in apply_primitive
compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)

File [...]/jax/interpreters/xla.py, line 262, in xla_primitive_callable
*avals, **params)

File [...]/jax/interpreters/xla.py, line 320, in primitive_computation
raise RuntimeError(msg) from e

RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.:
This is a bug in JAX's shape-checking rules; please report it!