beam_search.Search() in a single-GPU environment about trax HOT 7 CLOSED

TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 1) for shape (1, 2, 4, 1).

    seqs = seqs[:, :, 1:]  # Strip start token
IndexError: too many indices for array

[[[8 8 8 ... 8 8 8]]]
All done.

LayerError: Exception passing through layer Cache (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 270
  layer input shapes: (ShapeDtype{shape:(1, 2), dtype:int64}, ShapeDtype{shape:(1, 1, 1, 2), dtype:bool})

  File [...]/trax/layers/combinators.py, line 671, in forward_with_state
    if state[0] is ():  # pylint: disable=literal-comparison

IndexError: tuple index out of range

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/content/trax/trax/layers/base.py in pure_fn(self, x, weights, state, rng)
    446         outputs, s = (
--> 447             self.forward_with_state(x, weights=weights, state=state, rng=rng))
    448       else:

14 frames
/content/trax/trax/layers/combinators.py in forward_with_state(self, xs, weights, state, rng)
     67       raise ValueError(
---> 68           f'Number of weight elements ({len(weights)}) does not equal '
     69           f'number of sublayers ({n_layers}).')

ValueError: Number of weight elements (16) does not equal number of sublayers (31).

The above exception was the direct cause of the following exception:

LayerError                                Traceback (most recent call last)
<ipython-input-45-43120f5a1f8f> in <module>()
      1 
----> 2 preds, scores = sampling_decoder.decode(inputs)
      3 
      4 print(preds)
      5 print('All done.')

/content/trax/trax/models/beam_search.py in decode(self, inputs, targets_prefix, batch_size)
    593     seqs, scores = self._jit_beam_search(
    594         inputs, targets_prefix, (batch_size + pad_amount) // n_devices,
--> 595         dummy=np.zeros(n_devices))
    596     seqs = onp.asarray(seqs)
    597     scores = onp.asarray(scores)

/usr/local/lib/python3.6/dist-packages/jax/api.py in f_pmapped(*args, **kwargs)
    930         global_axis_size=axis_size,
    931         devices=tuple(devices) if devices is not None else devices,
--> 932         name=flat_fun.__name__)
    933     return tree_unflatten(out_tree(), out)
    934 

/usr/local/lib/python3.6/dist-packages/jax/core.py in call_bind(primitive, f, *args, **params)
    893   if top_trace is None:
    894     with new_sublevel():
--> 895       outs = primitive.impl(f, *args, **params)
    896   else:
    897     tracers = map(top_trace.full_raise, args)

/usr/local/lib/python3.6/dist-packages/jax/interpreters/pxla.py in xla_pmap_impl(fun, backend, axis_name, axis_size, global_axis_size, devices, name, mapped_invars, *args)
    404   abstract_args = map(xla.abstractify, args)
    405   compiled_fun = parallel_callable(fun, backend, axis_name, axis_size,
--> 406                                    global_axis_size, devices, name, *abstract_args)
    407   return compiled_fun(*args)
    408 

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in memoized_fun(fun, *args)
    218       fun.populate_stores(stores)
    219     else:
--> 220       ans = call(fun, *args)
    221       cache[key] = (ans, fun.stores)
    222     return ans

/usr/local/lib/python3.6/dist-packages/jax/interpreters/pxla.py in parallel_callable(fun, backend, axis_name, axis_size, global_axis_size, devices, name, *avals)
    460   pval = pe.PartialVal([core.abstract_unit, core.unit])  # dummy value for axis env
    461   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(
--> 462       dynamic_fun, [pval] + pvals, instantiate=False, stage_out_calls=True, bottom=True)
    463   jaxpr.invars = jaxpr.invars[1:]  # ignore dummy
    464 

/usr/local/lib/python3.6/dist-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out_calls, bottom)
    356   with new_master(trace_type, bottom=bottom) as master:
    357     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 358     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    359     assert not env
    360     del master

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    147     gen = None
    148 
--> 149     ans = self.f(*args, **dict(self.params, **kwargs))
    150     del args
    151     while stack:

/usr/local/lib/python3.6/dist-packages/jax/interpreters/pxla.py in dynamic_fun(dummy, *args)
    453   def dynamic_fun(dummy, *args):
    454     with extend_dynamic_axis_env(axis_name, dummy._trace, global_axis_size):
--> 455       return fun.call_wrapped(*args)
    456 
    457   avals = tuple(map(partial(shard_aval, axis_size), avals))

/usr/local/lib/python3.6/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    147     gen = None
    148 
--> 149     ans = self.f(*args, **dict(self.params, **kwargs))
    150     del args
    151     while stack:

/content/trax/trax/models/beam_search.py in _unreplicated_beam_search(***failed resolving arguments***)
    526     return beam_search(
    527         batch_size,
--> 528         self._get_initial_state(inputs, targets_prefix, batch_size),
    529         tokens_to_logits,
    530         max_decode_len,

/content/trax/trax/models/beam_search.py in _get_initial_state(self, inputs, targets_prefix, batch_size)
    487         self.model_weights,
    488         initial_state,
--> 489         jax.random.PRNGKey(0))
    490     state_structure = jax.tree_structure(prompted_state)
    491 

/content/trax/trax/layers/base.py in pure_fn(self, x, weights, state, rng)
    454       name, trace = self._name, _short_traceback()
    455       raise LayerError(name, 'pure_fn',
--> 456                        self._caller, signature(x), trace) from e
    457 
    458   def output_signature(self, input_signature):

LayerError: Exception passing through layer Serial (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 299
  layer input shapes: (ShapeDtype{shape:(1, 16), dtype:int32}, ShapeDtype{shape:(1, 1), dtype:int32})

  File [...]/trax/layers/combinators.py, line 68, in forward_with_state
    f'Number of weight elements ({len(weights)}) does not equal '

ValueError: Number of weight elements (16) does not equal number of sublayers (31).

# DO NOT FORGET TO PARSE GIN CONFIG WITH THE HYPERPARAMETERS YOU HAD DURING TRAINING
# otherwise, you will get layers error (number incopatibilities) and non-filled parameters of
# function calls 
model_infer_1 = trax.models.Transformer(mode='train')

model_infer_1.init_from_file(
    'gs://your/path/to/model.pkl', weights_only=False)

inputs = np.array([[vocab['<SOS>'], 1, 2, 3, 4, 5, 6, 7, 8, 9, vocab['<EOS>']]])

targets_prefix = np.array([[vocab['<SOS>'], 0]])

for char_i in range(inputs.shape[1]):
  ret = model_infer_1((inputs, targets_prefix))
  logits = ret[0]
  batch = logits[0]

  new_decoded_char = np.argmax(batch[-1], -1)

  if new_decoded_char == vocab['<EOS>']:
    output = targets_prefix[:, :-1]
    break

  # Take all previous chars except for the last one used to make a prediction,
  # then add the predicted one and new zero character to indicate that the new
  # character has to be searched
  targets_prefix = np.hstack((targets_prefix[:, :-1], np.array([[new_decoded_char, 0]],)))

  print(f'{targets_prefix}')