# Path: flaxformer/types.py



from typing import Any, Optional
from flax import linen as nn
from flax.linen.module import Module
from flax.training import common_utils
from jax import lax
from flaxformer.types import Array
import jax.numpy as jnp
  shift[axis] = distance
  return roll_with_zeros(x, shift)


class Depthwise1dConv(Module):
  """One-dimensional depthwise convolution.

  If autoregressive=True, then position `i` receives information from positions
  in the interval `[i-radius, i]`.

  If autoregressive=False, then position `i` receives information from positions
  in the interval `[i-radius, i+radius]`.

  Attributes:
    radius: Maximum distance to move information.
    autoregressive: Whether to only look left.
    dtype: the dtype of the computation (default: float32).
  """
  radius: int = 2
  autoregressive: bool = True
  dtype: Any = jnp.float32
  length_dim: int = 1
  num_feature_dims: int = 1
  use_in_mlp_parallel_fused_layer: bool = False

  @nn.compact
  def __call__(self,
               x: Array,
               decode: bool = False,
               prefill: bool = False,

# Path: flaxformer/types.py



# Path: flaxformer/components/convolution.py
from typing import Any, Optional
from flax import linen as nn
from flax.linen.module import Module
from flax.training import common_utils
from jax import lax
from flaxformer.types import Array
import jax.numpy as jnp

  shift[axis] = distance
  return roll_with_zeros(x, shift)


class Depthwise1dConv(Module):
  """One-dimensional depthwise convolution.

  If autoregressive=True, then position `i` receives information from positions
  in the interval `[i-radius, i]`.

  If autoregressive=False, then position `i` receives information from positions
  in the interval `[i-radius, i+radius]`.

  Attributes:
    radius: Maximum distance to move information.
    autoregressive: Whether to only look left.
    dtype: the dtype of the computation (default: float32).
  """
  radius: int = 2
  autoregressive: bool = True
  dtype: Any = jnp.float32
  length_dim: int = 1
  num_feature_dims: int = 1
  use_in_mlp_parallel_fused_layer: bool = False

  @nn.compact
  def __call__(self,
               x: Array,
               decode: bool = False,
               prefill: bool = False,

               prefill_lengths: Optional[Array] = None) -> Array: