Repository: google/trax

Issue:
beam_search.Search() in a single-GPU environment
### Description
Training a Transformer converges. 

Then beam_search fails though. When n_devices == 1 some reshapes crash  in decode().

### Environment information

```
OS: 
ubuntu 18.04 

CUDA 10.1
1 GPU environment

$ pip freeze | grep tensor
mesh-tensorflow==0.1.11
tensor2tensor==1.15.4
tensorboard==2.1.0
tensorflow==2.1.0
tensorflow-datasets==2.1.0
tensorflow-estimator==2.1.0
tensorflow-gan==2.0.0
tensorflow-hub==0.7.0
tensorflow-metadata==0.21.1
tensorflow-probability==0.7.0


$ pip freeze | grep jax
jax==0.1.59
jaxlib==0.1.39


$ python -V
Python 3.6.9
```

### For bugs: reproduction and error logs
Traceback (most recent call last):
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 442, in pure_fn
    x, weights=weights, state=state, rng=rng)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 220, in forward_with_state
    return self.forward(inputs, weights), state
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 580, in _forward
    raw_output = raw_fn(x, weights=weights, **self._kwargs)  # pylint: disable=protected-access
  File "/root/.local/lib/python3.6/site-packages/trax/layers/attention.py", line 51, in PaddingMask
    return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))
  File "/root/.local/lib/python3.6/site-packages/jax/numpy/lax_numpy.py", line 921, in reshape
    return a.reshape(newshape, order=order)  # forward to method for ndarrays
  File "/root/.local/lib/python3.6/site-packages/jax/numpy/lax_numpy.py", line 959, in _reshape_method
    return _reshape(a, newshape, order=order)
  File "/root/.local/lib/python3.6/site-packages/jax/numpy/lax_numpy.py", line 938, in _reshape
    return lax.reshape(a, computed_newshape, None)
  File "/root/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 640, in reshape
    old_sizes=onp.shape(operand))
  File "/root/.local/lib/python3.6/site-packages/jax/core.py", line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/root/.local/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 98, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)
  File "/root/.local/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 106, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)
  File "/root/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 1523, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
  File "/root/.local/lib/python3.6/site-packages/jax/lax/lax.py", line 2582, in _reshape_shape_rule
    raise TypeError(msg.format(new_sizes, onp.shape(operand)))
TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 30) for shape (1, 2, 30).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 480, in _forward_abstract
    input_signature, weight_signature, self.state, rng)
  File "/root/.local/lib/python3.6/site-packages/trax/math/jax.py", line 175, in shape_fun
    jax_shapes = jax.eval_shape(f, *args, **kwargs)
  File "/root/.local/lib/python3.6/site-packages/jax/api.py", line 2042, in eval_shape
    out = pe.abstract_eval_fun(fun.call_wrapped, *map(abstractify, args_flat))
  File "/root/.local/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 273, in abstract_eval_fun
    instantiate=True)
  File "/root/.local/lib/python3.6/site-packages/jax/interpreters/partial_eval.py", line 354, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/root/.local/lib/python3.6/site-packages/jax/linear_util.py", line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/root/.local/lib/python3.6/site-packages/jax/linear_util.py", line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 477, in call_on_input
    return self.forward_with_state(x, weights=weights, state=state, rng=rng)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/combinators.py", line 238, in forward_with_state
    sub_outputs, sub_state = layer.pure_fn(x, w, s, r)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 451, in pure_fn
    self._caller, signature(x), trace)
trax.layers.base.LayerError: Exception passing through layer PaddingMask (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 286
  layer input shapes: ShapeDtype{shape:(1, 2, 30), dtype:int32}

  File [...]/trax/layers/base.py, line 220, in forward_with_state
    return self.forward(inputs, weights), state

  File [...]/trax/layers/base.py, line 580, in _forward
    raw_output = raw_fn(x, weights=weights, **self._kwargs)  # pylint: disable=protected-access

  File [...]/trax/layers/attention.py, line 51, in PaddingMask
    return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))

  File [...]/jax/numpy/lax_numpy.py, line 921, in reshape
    return a.reshape(newshape, order=order)  # forward to method for ndarrays

  File [...]/jax/numpy/lax_numpy.py, line 959, in _reshape_method
    return _reshape(a, newshape, order=order)

  File [...]/jax/numpy/lax_numpy.py, line 938, in _reshape
    return lax.reshape(a, computed_newshape, None)

  File [...]/jax/lax/lax.py, line 640, in reshape
    old_sizes=onp.shape(operand))

  File [...]/site-packages/jax/core.py, line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)

  File [...]/jax/interpreters/partial_eval.py, line 98, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)

  File [...]/jax/interpreters/partial_eval.py, line 106, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)

  File [...]/jax/lax/lax.py, line 1523, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))

  File [...]/jax/lax/lax.py, line 2582, in _reshape_shape_rule
    raise TypeError(msg.format(new_sizes, onp.shape(operand)))

TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 30) for shape (1, 2, 30).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 310, in init
    weights, state = self.new_weights_and_state(input_signature)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/combinators.py", line 92, in new_weights_and_state
    outputs, _ = sublayer._forward_abstract(inputs)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 485, in _forward_abstract
    trace)
trax.layers.base.LayerError: Exception passing through layer Parallel (in _forward_abstract):
  layer created in file [...]/trax/layers/combinators.py, line 468
  layer input shapes: (ShapeDtype{shape:(1, 2, 30), dtype:int32}, ShapeDtype{shape:(1, 2, 30), dtype:int32})

  File [...]/trax/math/jax.py, line 175, in shape_fun
    jax_shapes = jax.eval_shape(f, *args, **kwargs)

  File [...]/site-packages/jax/api.py, line 2042, in eval_shape
    out = pe.abstract_eval_fun(fun.call_wrapped, *map(abstractify, args_flat))

  File [...]/jax/interpreters/partial_eval.py, line 273, in abstract_eval_fun
    instantiate=True)

  File [...]/jax/interpreters/partial_eval.py, line 354, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)

  File [...]/site-packages/jax/linear_util.py, line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/site-packages/jax/linear_util.py, line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/trax/layers/base.py, line 477, in call_on_input
    return self.forward_with_state(x, weights=weights, state=state, rng=rng)

  File [...]/trax/layers/combinators.py, line 238, in forward_with_state
    sub_outputs, sub_state = layer.pure_fn(x, w, s, r)

LayerError: Exception passing through layer PaddingMask (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 286
  layer input shapes: ShapeDtype{shape:(1, 2, 30), dtype:int32}

  File [...]/trax/layers/base.py, line 220, in forward_with_state
    return self.forward(inputs, weights), state

  File [...]/trax/layers/base.py, line 580, in _forward
    raw_output = raw_fn(x, weights=weights, **self._kwargs)  # pylint: disable=protected-access

  File [...]/trax/layers/attention.py, line 51, in PaddingMask
    return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))

  File [...]/jax/numpy/lax_numpy.py, line 921, in reshape
    return a.reshape(newshape, order=order)  # forward to method for ndarrays

  File [...]/jax/numpy/lax_numpy.py, line 959, in _reshape_method
    return _reshape(a, newshape, order=order)

  File [...]/jax/numpy/lax_numpy.py, line 938, in _reshape
    return lax.reshape(a, computed_newshape, None)

  File [...]/jax/lax/lax.py, line 640, in reshape
    old_sizes=onp.shape(operand))

  File [...]/site-packages/jax/core.py, line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)

  File [...]/jax/interpreters/partial_eval.py, line 98, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)

  File [...]/jax/interpreters/partial_eval.py, line 106, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)

  File [...]/jax/lax/lax.py, line 1523, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))

  File [...]/jax/lax/lax.py, line 2582, in _reshape_shape_rule
    raise TypeError(msg.format(new_sizes, onp.shape(operand)))

TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 30) for shape (1, 2, 30).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 310, in init
    weights, state = self.new_weights_and_state(input_signature)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/combinators.py", line 91, in new_weights_and_state
    weights_or_empty, state = sublayer.init(inputs)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 321, in init
    input_signature, trace)
trax.layers.base.LayerError: Exception passing through layer Serial (in init):
  layer created in file [...]/trax/layers/combinators.py, line 470
  layer input shapes: ShapeDtype{shape:(1, 2, 30), dtype:int32}

  File [...]/trax/layers/combinators.py, line 92, in new_weights_and_state
    outputs, _ = sublayer._forward_abstract(inputs)

LayerError: Exception passing through layer Parallel (in _forward_abstract):
  layer created in file [...]/trax/layers/combinators.py, line 468
  layer input shapes: (ShapeDtype{shape:(1, 2, 30), dtype:int32}, ShapeDtype{shape:(1, 2, 30), dtype:int32})

  File [...]/trax/math/jax.py, line 175, in shape_fun
    jax_shapes = jax.eval_shape(f, *args, **kwargs)

  File [...]/site-packages/jax/api.py, line 2042, in eval_shape
    out = pe.abstract_eval_fun(fun.call_wrapped, *map(abstractify, args_flat))

  File [...]/jax/interpreters/partial_eval.py, line 273, in abstract_eval_fun
    instantiate=True)

  File [...]/jax/interpreters/partial_eval.py, line 354, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)

  File [...]/site-packages/jax/linear_util.py, line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/site-packages/jax/linear_util.py, line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/trax/layers/base.py, line 477, in call_on_input
    return self.forward_with_state(x, weights=weights, state=state, rng=rng)

  File [...]/trax/layers/combinators.py, line 238, in forward_with_state
    sub_outputs, sub_state = layer.pure_fn(x, w, s, r)

LayerError: Exception passing through layer PaddingMask (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 286
  layer input shapes: ShapeDtype{shape:(1, 2, 30), dtype:int32}

  File [...]/trax/layers/base.py, line 220, in forward_with_state
    return self.forward(inputs, weights), state

  File [...]/trax/layers/base.py, line 580, in _forward
    raw_output = raw_fn(x, weights=weights, **self._kwargs)  # pylint: disable=protected-access

  File [...]/trax/layers/attention.py, line 51, in PaddingMask
    return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))

  File [...]/jax/numpy/lax_numpy.py, line 921, in reshape
    return a.reshape(newshape, order=order)  # forward to method for ndarrays

  File [...]/jax/numpy/lax_numpy.py, line 959, in _reshape_method
    return _reshape(a, newshape, order=order)

  File [...]/jax/numpy/lax_numpy.py, line 938, in _reshape
    return lax.reshape(a, computed_newshape, None)

  File [...]/jax/lax/lax.py, line 640, in reshape
    old_sizes=onp.shape(operand))

  File [...]/site-packages/jax/core.py, line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)

  File [...]/jax/interpreters/partial_eval.py, line 98, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)

  File [...]/jax/interpreters/partial_eval.py, line 106, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)

  File [...]/jax/lax/lax.py, line 1523, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))

  File [...]/jax/lax/lax.py, line 2582, in _reshape_shape_rule
    raise TypeError(msg.format(new_sizes, onp.shape(operand)))

TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 30) for shape (1, 2, 30).

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "math_trax.py", line 565, in <module>
    seqs, scores = beam_decoder.decode(inputs=batch, batch_size=iBatch_size)#, )targets_prefix=prefix_for_bs,
  File "/root/.local/lib/python3.6/site-packages/trax/models/beam_search.py", line 602, in decode
    dummy=np.zeros(n_devices))
  File "/root/.local/lib/python3.6/site-packages/jax/api.py", line 146, in f_jitted
    name=flat_fun.__name__)
  File "/root/.local/lib/python3.6/site-packages/jax/core.py", line 642, in call_bind
    outs = primitive.impl(f, *args, **params)
  File "/root/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 448, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, *map(arg_spec, args))
  File "/root/.local/lib/python3.6/site-packages/jax/linear_util.py", line 220, in memoized_fun
    ans = call(fun, *args)
  File "/root/.local/lib/python3.6/site-packages/jax/interpreters/xla.py", line 465, in _xla_callable
    jaxpr, (pvals, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
  File "/root/.local/lib/python3.6/site-packages/jax/linear_util.py", line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/root/.local/lib/python3.6/site-packages/trax/models/beam_search.py", line 535, in _unreplicated_beam_search
    self._get_initial_state(inputs, targets_prefix, batch_size),
  File "/root/.local/lib/python3.6/site-packages/trax/models/beam_search.py", line 490, in _get_initial_state
    _, initial_state = self.model(mode='predict').init(signature)
  File "/root/.local/lib/python3.6/site-packages/trax/layers/base.py", line 321, in init
    input_signature, trace)
trax.layers.base.LayerError: Exception passing through layer Serial (in init):
  layer created in file [...]/trax/models/transformer.py, line 301
  layer input shapes: (ShapeDtype{shape:(1, 2, 30), dtype:int32}, ShapeDtype{shape:(2, 1), dtype:int32})

  File [...]/trax/layers/combinators.py, line 91, in new_weights_and_state
    weights_or_empty, state = sublayer.init(inputs)

LayerError: Exception passing through layer Serial (in init):
  layer created in file [...]/trax/layers/combinators.py, line 470
  layer input shapes: ShapeDtype{shape:(1, 2, 30), dtype:int32}

  File [...]/trax/layers/combinators.py, line 92, in new_weights_and_state
    outputs, _ = sublayer._forward_abstract(inputs)

LayerError: Exception passing through layer Parallel (in _forward_abstract):
  layer created in file [...]/trax/layers/combinators.py, line 468
  layer input shapes: (ShapeDtype{shape:(1, 2, 30), dtype:int32}, ShapeDtype{shape:(1, 2, 30), dtype:int32})

  File [...]/trax/math/jax.py, line 175, in shape_fun
    jax_shapes = jax.eval_shape(f, *args, **kwargs)

  File [...]/site-packages/jax/api.py, line 2042, in eval_shape
    out = pe.abstract_eval_fun(fun.call_wrapped, *map(abstractify, args_flat))

  File [...]/jax/interpreters/partial_eval.py, line 273, in abstract_eval_fun
    instantiate=True)

  File [...]/jax/interpreters/partial_eval.py, line 354, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)

  File [...]/site-packages/jax/linear_util.py, line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/site-packages/jax/linear_util.py, line 149, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/trax/layers/base.py, line 477, in call_on_input
    return self.forward_with_state(x, weights=weights, state=state, rng=rng)

  File [...]/trax/layers/combinators.py, line 238, in forward_with_state
    sub_outputs, sub_state = layer.pure_fn(x, w, s, r)

LayerError: Exception passing through layer PaddingMask (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 286
  layer input shapes: ShapeDtype{shape:(1, 2, 30), dtype:int32}

  File [...]/trax/layers/base.py, line 220, in forward_with_state
    return self.forward(inputs, weights), state

  File [...]/trax/layers/base.py, line 580, in _forward
    raw_output = raw_fn(x, weights=weights, **self._kwargs)  # pylint: disable=protected-access

  File [...]/trax/layers/attention.py, line 51, in PaddingMask
    return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))

  File [...]/jax/numpy/lax_numpy.py, line 921, in reshape
    return a.reshape(newshape, order=order)  # forward to method for ndarrays

  File [...]/jax/numpy/lax_numpy.py, line 959, in _reshape_method
    return _reshape(a, newshape, order=order)

  File [...]/jax/numpy/lax_numpy.py, line 938, in _reshape
    return lax.reshape(a, computed_newshape, None)

  File [...]/jax/lax/lax.py, line 640, in reshape
    old_sizes=onp.shape(operand))

  File [...]/site-packages/jax/core.py, line 182, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)

  File [...]/jax/interpreters/partial_eval.py, line 98, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)

  File [...]/jax/interpreters/partial_eval.py, line 106, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)

  File [...]/jax/lax/lax.py, line 1523, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))

  File [...]/jax/lax/lax.py, line 2582, in _reshape_shape_rule
    raise TypeError(msg.format(new_sizes, onp.shape(operand)))

TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 30) for shape (1, 2, 30).

REMARK:
1 is n_devices, 2 is batch size, 30 is max_len


```
# Steps to reproduce:
I tried to force your machine_translation.ipynb in colab to use the GPU but didnt succeed. But maybe for you it's the fastest to check what happens if only 1 GPU as the colab in itsef runs smoothly (on a TPU).
```

```
# Error logs:
...
```



Discussion:
Hmmm. Possibly this is not an n_devices-in-decode-issue but a bug in Transformer() that appears in beam-search-settings as I can replicate the bug from my local system (GPU only) also in your colab if I upload a Tranformer-model (rather than a Reformer-model)
@friesel I am getting the same now:


Running:

```python
import os
import trax
import gin
import numpy as np
from tensorflow_datasets.core.features.text import SubwordTextEncoder
from trax.models.beam_search import Search


def main():

    output_dir = '/tmp/trax/model'
    vocab_fp = os.path.join(output_dir, 'vocab')

    text_encoder = SubwordTextEncoder.load_from_file(vocab_fp)

    gin.parse_config_file('examples/transformer.cfg')
    gin.bind_parameter('Transformer.input_vocab_size', text_encoder.vocab_size)

    model_infer = trax.models.Transformer(mode='predict')
    model_infer.init_from_file('/tmp/trax/model/model.pkl', weights_only=True)

    sampling_decoder = Search(
        trax.models.Transformer,
        model_infer.weights,
        temperature=1.0,
        max_decode_len=32 * 64 * 3,
    )

    inputs = np.asarray(
        [
            [4, 5, 6, 1],
            [5, 7, 8, 1]
        ]
    )

    preds, scores = sampling_decoder.decode(inputs)

    print('All done.')


if __name__ == '__main__':
    main()
```

Will throw 

```none
TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 1) for shape (1, 2, 4, 1).
```


<details>
   <summary>Click to expand entire error message</summary>

```none
trax.layers.base.LayerError: Exception passing through layer Serial (in init):
  layer created in file [...]/trax/models/transformer.py, line 299
  layer input shapes: (ShapeDtype{shape:(1, 2, 4, 1), dtype:int32}, ShapeDtype{shape:(2, 1), dtype:int32})

  File [...]/trax/layers/combinators.py, line 91, in new_weights_and_state
    weights_or_empty, state = sublayer.init(inputs)

LayerError: Exception passing through layer Branch (in init):
  layer created in file [...]/trax/layers/combinators.py, line 472
  layer input shapes: ShapeDtype{shape:(1, 2, 4, 1), dtype:int32}

  File [...]/trax/layers/combinators.py, line 92, in new_weights_and_state
    outputs, _ = sublayer._forward_abstract(inputs)

LayerError: Exception passing through layer Parallel (in _forward_abstract):
  layer created in file [...]/trax/layers/combinators.py, line 470
  layer input shapes: (ShapeDtype{shape:(1, 2, 4, 1), dtype:int32}, ShapeDtype{shape:(1, 2, 4, 1), dtype:int32})

  File [...]/trax/math/jax.py, line 175, in shape_fun
    jax_shapes = jax.eval_shape(f, *args, **kwargs)

  File [...]/site-packages/jax/api.py, line 1602, in eval_shape
    *map(abstractify, args_flat))

  File [...]/jax/interpreters/partial_eval.py, line 326, in abstract_eval_fun
    instantiate=True, stage_out=True)

  File [...]/jax/interpreters/partial_eval.py, line 421, in trace_to_jaxpr
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)

  File [...]/site-packages/jax/linear_util.py, line 150, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/site-packages/jax/linear_util.py, line 150, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))

  File [...]/trax/layers/base.py, line 479, in call_on_input
    return self.forward_with_state(x, weights=weights, state=state, rng=rng)

  File [...]/trax/layers/combinators.py, line 239, in forward_with_state
    sub_outputs, sub_state = layer.pure_fn(x, w, s, r)

LayerError: Exception passing through layer PaddingMask (in pure_fn):
  layer created in file [...]/trax/models/transformer.py, line 284
  layer input shapes: ShapeDtype{shape:(1, 2, 4, 1), dtype:int32}

  File [...]/trax/layers/base.py, line 222, in forward_with_state
    return self.forward(inputs, weights), state

  File [...]/trax/layers/base.py, line 582, in _forward
    raw_output = raw_fn(x, weights=weights, **self._kwargs)  # pylint: disable=protected-access

  File [...]/trax/layers/attention.py, line 55, in PaddingMask
    return np.reshape(x != pad, (x.shape[0], 1, 1, x.shape[-1]))

  File [...]/jax/numpy/lax_numpy.py, line 1016, in reshape
    return a.reshape(newshape, order=order)  # forward to method for ndarrays

  File [...]/jax/numpy/lax_numpy.py, line 1054, in _reshape_method
    return _reshape(a, newshape, order=order)

  File [...]/jax/numpy/lax_numpy.py, line 1033, in _reshape
    return lax.reshape(a, computed_newshape, None)

  File [...]/jax/lax/lax.py, line 675, in reshape
    dimensions=None if dimensions is None or same_dims else tuple(dimensions))

  File [...]/site-packages/jax/core.py, line 202, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)

  File [...]/jax/interpreters/partial_eval.py, line 133, in process_primitive
    return self.default_process_primitive(primitive, tracers, params)

  File [...]/jax/interpreters/partial_eval.py, line 141, in default_process_primitive
    out_aval = primitive.abstract_eval(*avals, **params)

  File [...]/jax/lax/lax.py, line 1672, in standard_abstract_eval
    return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))

  File [...]/jax/lax/lax.py, line 2856, in _reshape_shape_rule
    raise TypeError(msg.format(new_sizes, onp.shape(operand)))

TypeError: reshape total size must be unchanged, got new_sizes (1, 1, 1, 1) for shape (1, 2, 4, 1).

Process finished with exit code 1
```
</details>

However, I am able to get passt this error when presenting an input like this:

```python
    inputs = np.asarray(
        [[
            [4, 5, 6, 1]
        ]]
    )
```

In this case I get predictions but there will be another crash right after that in 

https://github.com/google/trax/blob/5b1565910a53d0d1175f647cc67db48e334d8f90/trax/models/beam_search.py#L602

![image](https://user-images.githubusercontent.com/43335432/79771805-94b7da80-832f-11ea-94de-462acb4874be.png)

However, those predictions probably not what I want anyway. 

**Update**

I noticied that `Search` has an argument `eos_id`. 

```python
    sampling_decoder = Search(
        trax.models.Transformer,
        model_infer.weights,
        temperature=1.0,
        max_decode_len=32 * 64 * 3,
        eos_id=EOS_ID
    )
```

Setting this to the actual one that I am using, `1` instead of `-1` (default), I get through the call of `self._jit_beam_search` but then end up with another error that is caused by




https://github.com/google/trax/blob/5b1565910a53d0d1175f647cc67db48e334d8f90/trax/models/beam_search.py#L602

because `seq`, at this point, is two dimensional. Leaving me at:

```
    seqs = seqs[:, :, 1:]  # Strip start token
IndexError: too many indices for array
```


Resolving patch:
diff --git a/trax/models/beam_search.py b/trax/models/beam_search.py
--- a/trax/models/beam_search.py
+++ b/trax/models/beam_search.py
@@ -459,11 +459,8 @@ def __init__(self, model, weights, max_decode_len, beam_size=1, temperature=0,
         eos_token=eos_id if eos_id is not None else -1,
         max_decode_len=max_decode_len + 1)  # Add 1 to account for start token.
 
-    if trax.math.device_count() == 1:
-      self._jit_beam_search = jax.jit(beam_search_partial, static_argnums=(2,))
-    else:
-      self._jit_beam_search = jax.pmap(beam_search_partial, axis_name='batch',
-                                       static_broadcasted_argnums=(2,))
+    self._jit_beam_search = jax.pmap(beam_search_partial, axis_name='batch',
+                                     static_broadcasted_argnums=(2,))
 
   def _get_initial_state(self, inputs, targets_prefix, batch_size):
     """Get initial state for beam search."""
@@ -511,8 +508,9 @@ def _unreplicated_beam_search(self, inputs, targets_prefix, batch_size, dummy,
     def tokens_to_logits(flat_ids, flat_cache, rng):
       """Autoregressive decoding step: map from previous tokens to logits."""
       rng, subrng = jax.random.split(rng)
+      tiled_inputs = flatten_beam_dim(add_beam_dim(inputs, beam_size))
       flat_logits, new_flat_cache = self.model_infer.pure_fn(
-          flat_ids if inputs is None else (inputs, flat_ids),
+          flat_ids if tiled_inputs is None else (tiled_inputs, flat_ids),
           self.model_weights,
           flat_cache,
           rng)
diff --git a/trax/models/reformer/reformer.py b/trax/models/reformer/reformer.py
--- a/trax/models/reformer/reformer.py
+++ b/trax/models/reformer/reformer.py
@@ -1127,8 +1127,8 @@ def PositionalEncoder(vocab_size, mode):  # tokens --> vectors
   return tl.Serial(
       # Input: encoder_side_tokens, decoder_side_tokens
       # Copy decoder tokens for use in loss.
-      tl.Branch([], [tl.Dup(),              # tok_e tok_d tok_d
-                     tl.PaddingMask(),
+      tl.Select([0, 1, 1]),                 # tok_e tok_d tok_d
+      tl.Branch([], [tl.PaddingMask(),
                      tl.Fn(lambda x: np.squeeze(x, (1, 2)), n_out=1)]),
       #                                     # tok_e mask  tok_d .....
 


Test patch: