@tf.keras.utils.register_keras_serializable(package="lite_mamba")
class TFMamba(tf.keras.layers.Layer):
"""TensorFlow Mamba block with parallel dilated depthwise causal conv branches."""
def __init__(
self,
d_model,
d_state=16,
d_conv=4,
conv_dilations=(1,),
expand=2,
dt_rank="auto",
dt_min=0.001,
dt_max=0.1,
dt_init="random",
dt_scale=1.0,
dt_init_floor=1e-4,
conv_bias=True,
bias=False,
layer_idx=None,
stacked_convs=False,
pointwise=False,
**kwargs,
):
super().__init__(**kwargs)
if dt_init not in ("constant", "random"):
raise ValueError(f"dt_init must be 'constant' or 'random', got {dt_init!r}.")
self.d_model = d_model
self.d_state = d_state
self.d_conv = d_conv
self.conv_dilations = tuple(conv_dilations)
self.num_conv_branches = len(self.conv_dilations)
self.conv_state_lens = [(self.d_conv - 1) * d + 1 for d in self.conv_dilations]
self.expand = expand
self.d_inner = int(self.expand * self.d_model)
self.dt_rank = math.ceil(self.d_model / 16) if dt_rank == "auto" else dt_rank
self.dt_min = dt_min
self.dt_max = dt_max
self.dt_init = dt_init
self.dt_scale = dt_scale
self.dt_init_floor = dt_init_floor
self.conv_bias = conv_bias
self.bias = bias
self.layer_idx = layer_idx
self.stacked_convs = stacked_convs
self.pointwise = pointwise
# `dt_proj` is only ever used as a weight container — `call`/`step` read its
# kernel and bias directly so that the bias can be folded into the scan's
# softplus. Both are initialised declaratively here, which keeps `build()`
# free of post-hoc `assign()` calls and therefore idempotent.
dt_init_std = (self.dt_rank**-0.5) * self.dt_scale
dt_kernel_init = (
tf.keras.initializers.Constant(dt_init_std)
if self.dt_init == "constant"
else tf.keras.initializers.RandomUniform(-dt_init_std, dt_init_std)
)
self.in_proj = tf.keras.layers.Dense(self.d_inner * 2, use_bias=self.bias, name="in_proj")
self.x_proj = tf.keras.layers.Dense(
self.dt_rank + self.d_state * 2, use_bias=False, name="x_proj"
)
self.dt_proj = tf.keras.layers.Dense(
self.d_inner,
use_bias=True,
kernel_initializer=dt_kernel_init,
bias_initializer=InverseSoftplusDtInitializer(
self.dt_min, self.dt_max, self.dt_init_floor
),
name="dt_proj",
)
self.out_proj = tf.keras.layers.Dense(self.d_model, use_bias=self.bias, name="out_proj")
self.pw_layers = (
[
tf.keras.layers.Dense(self.d_inner, use_bias=True, name=f"pw_proj_{i}")
for i in range(self.num_conv_branches)
]
if self.pointwise
else []
)
# Weights are created in `build()` (not `__init__`) so that they inherit the
# layer's name scope. Creating them here would emit bare names such as
# `A_log:0`, and — combined with the sub-layer builds below — produced the
# duplicate-weight-name collisions that broke HDF5 serialisation.
self.conv_gates = None
self.dw_kernels = []
self.dw_biases = []
self.A_log = None
self.D = None
def build(self, input_shape):
if self.built:
return
# `rank` is None for an unknown-rank shape, so this only fires when the
# feature dimension is actually known and wrong.
shape = tf.TensorShape(input_shape) if input_shape is not None else None
if shape is not None and shape.rank:
in_dim = shape[-1]
if in_dim is not None and int(in_dim) != self.d_model:
raise ValueError(
f"{type(self).__name__} was configured with d_model={self.d_model} "
f"but received input with last dimension {int(in_dim)}."
)
if self.num_conv_branches > 1 and not self.stacked_convs:
self.conv_gates = self.add_weight(
name="conv_gates",
shape=(self.num_conv_branches,),
initializer=tf.keras.initializers.Ones(),
trainable=True,
)
else:
self.conv_gates = None
self.dw_kernels = []
self.dw_biases = []
for i in range(self.num_conv_branches):
self.dw_kernels.append(
self.add_weight(
name=f"dw_kernel_{i}",
shape=(self.d_conv, self.d_inner),
initializer=tf.keras.initializers.GlorotUniform(),
trainable=True,
)
)
if self.conv_bias:
self.dw_biases.append(
self.add_weight(
name=f"dw_bias_{i}",
shape=(self.d_inner,),
initializer=tf.keras.initializers.Zeros(),
trainable=True,
)
)
else:
self.dw_biases.append(None)
self.A_log = self.add_weight(
name="A_log",
shape=(self.d_inner, self.d_state),
initializer=ALogInitializer(),
trainable=True,
)
self.D = self.add_weight(
name="D",
shape=(self.d_inner,),
initializer=tf.keras.initializers.Ones(),
trainable=True,
)
# Each sub-layer is built inside its own `tf.name_scope`. Calling
# `sub_layer.build(...)` directly bypasses `Layer.__call__`, which is what
# normally pushes the sub-layer's name scope — without this every Dense
# produced a variable literally named `<parent>/kernel:0`, so a single
# TFMamba block emitted four identically-named weights and
# `model.save_weights()` failed with "name already exists".
sublayers = [
(self.in_proj, self.d_model),
(self.x_proj, self.d_inner),
(self.dt_proj, self.dt_rank),
(self.out_proj, self.d_inner),
]
sublayers += [(pw, self.d_inner) for pw in self.pw_layers]
for layer, fan_in in sublayers:
with tf.name_scope(layer.name):
layer.build((None, None, fan_in))
super().build(input_shape)
def get_config(self):
config = super().get_config()
config.update(
{
"d_model": self.d_model,
"d_state": self.d_state,
"d_conv": self.d_conv,
"conv_dilations": self.conv_dilations,
"expand": self.expand,
"dt_rank": self.dt_rank,
"dt_min": self.dt_min,
"dt_max": self.dt_max,
"dt_init": self.dt_init,
"dt_scale": self.dt_scale,
"dt_init_floor": self.dt_init_floor,
"conv_bias": self.conv_bias,
"bias": self.bias,
"layer_idx": self.layer_idx,
"stacked_convs": self.stacked_convs,
"pointwise": self.pointwise,
}
)
return config
def compute_output_shape(self, input_shape):
return tuple(input_shape[:-1]) + (self.d_model,)
@staticmethod
def _causal_depthwise_conv1d(x, kernel, bias=None, dilation=1):
# x: (B, L, D), kernel: (K, D)
k = kernel.shape[0]
pad_left = dilation * (k - 1)
xpad = tf.pad(x, [[0, 0], [pad_left, 0], [0, 0]])
seqlen = tf.shape(x)[1]
taps = [xpad[:, i * dilation : i * dilation + seqlen, :] for i in range(k)]
stacked = tf.stack(taps, axis=2) # (B, L, K, D)
kernel = tf.cast(tf.reshape(kernel, [1, 1, k, -1]), stacked.dtype)
y = tf.reduce_sum(stacked * kernel, axis=2)
if bias is not None:
y = y + tf.cast(tf.reshape(bias, [1, 1, -1]), y.dtype)
return tf.nn.silu(y)
def _gate_weights(self, dtype):
"""Softmax over the branch gates, computed in float32 for stability."""
return tf.cast(tf.nn.softmax(tf.cast(self.conv_gates, tf.float32), axis=0), dtype)
def _run_branch(self, x, branch_idx):
y = self._causal_depthwise_conv1d(
x,
kernel=self.dw_kernels[branch_idx],
bias=self.dw_biases[branch_idx],
dilation=self.conv_dilations[branch_idx],
)
if self.pointwise:
y = self.pw_layers[branch_idx](y)
return y
def call(self, hidden_states, training=None):
# hidden_states: (B, L, D)
xz = self.in_proj(hidden_states)
x, z = tf.split(xz, num_or_size_splits=2, axis=-1)
A = -tf.exp(tf.cast(self.A_log, tf.float32))
if self.stacked_convs:
for i in range(self.num_conv_branches):
x = self._run_branch(x, i)
else:
conv_outputs = [self._run_branch(x, i) for i in range(self.num_conv_branches)]
if self.conv_gates is None:
x = conv_outputs[0]
else:
gate = self._gate_weights(conv_outputs[0].dtype)
x = tf.add_n([gate[i] * conv_outputs[i] for i in range(self.num_conv_branches)])
x_dbl = self.x_proj(x)
dt, B, C = tf.split(
x_dbl,
num_or_size_splits=[self.dt_rank, self.d_state, self.d_state],
axis=-1,
)
# Match PyTorch path: apply dt_proj weights here, then add bias/softplus inside scan.
dt = tf.einsum("blr,rd->bld", dt, tf.cast(self.dt_proj.kernel, dt.dtype))
y = selective_scan_tf(
x,
dt,
A,
B,
C,
self.D,
z=z,
delta_bias=self.dt_proj.bias,
delta_softplus=True,
return_last_state=False,
)
return self.out_proj(y)
def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None):
"""Allocate the recurrent caches consumed by :meth:`step`.
``dtype=None`` (the default) puts the conv cache in the layer's *compute*
dtype and the SSM state in float32. Hard-coding float32 here previously
broke ``step()`` under a ``mixed_float16`` policy, where activations arrive
as float16.
"""
conv_dtype = self.compute_dtype if dtype is None else dtype
ssm_dtype = tf.float32 if dtype is None else dtype
conv_state = [
tf.zeros([batch_size, state_len, self.d_inner], dtype=conv_dtype)
for state_len in self.conv_state_lens
]
ssm_state = tf.zeros([batch_size, self.d_inner, self.d_state], dtype=ssm_dtype)
return conv_state, ssm_state
def step(self, hidden_states, conv_state, ssm_state):
# hidden_states: (B, 1, D)
xz = self.in_proj(hidden_states[:, 0, :])
x, z = tf.split(xz, num_or_size_splits=2, axis=-1)
A = -tf.exp(tf.cast(self.A_log, tf.float32))
# `step` is invoked directly rather than through `Layer.__call__`, so Keras'
# autocast scope is not active and every variable read below yields the
# *variable* dtype. Casts here are therefore load-bearing under
# mixed_float16, not defensive.
compute_dtype = x.dtype
def branch_step(x_in, state, kernel, bias, dilation, pw_layer):
state = tf.cast(state, x_in.dtype)
x_state = tf.concat([state[:, 1:, :], tf.expand_dims(x_in, axis=1)], axis=1)
k = kernel.shape[0]
idx = tf.range(k - 1, -1, -1) * dilation
pos = tf.shape(x_state)[1] - 1 - idx
values = tf.gather(x_state, pos, axis=1) # (B, K, D), oldest->newest
weights = tf.cast(tf.reshape(kernel, [1, k, self.d_inner]), values.dtype)
y = tf.reduce_sum(values * weights, axis=1)
if bias is not None:
y = y + tf.cast(bias, y.dtype)
y = tf.nn.silu(y)
if pw_layer is not None:
y = pw_layer(y)
return y, x_state
if self.stacked_convs:
new_states = []
for i in range(self.num_conv_branches):
pw = self.pw_layers[i] if self.pointwise else None
x, new_state = branch_step(
x,
conv_state[i],
self.dw_kernels[i],
self.dw_biases[i],
self.conv_dilations[i],
pw,
)
new_states.append(new_state)
else:
branch_outputs = []
new_states = []
for i in range(self.num_conv_branches):
pw = self.pw_layers[i] if self.pointwise else None
xi, new_state = branch_step(
x,
conv_state[i],
self.dw_kernels[i],
self.dw_biases[i],
self.conv_dilations[i],
pw,
)
branch_outputs.append(xi)
new_states.append(new_state)
if self.conv_gates is None:
x = branch_outputs[0]
else:
gate = self._gate_weights(branch_outputs[0].dtype)
x = tf.add_n([gate[i] * branch_outputs[i] for i in range(self.num_conv_branches)])
x_db = self.x_proj(x)
dt, B, C = tf.split(
x_db,
num_or_size_splits=[self.dt_rank, self.d_state, self.d_state],
axis=-1,
)
# The recurrence runs in float32 to mirror `selective_scan_tf`, so that
# `step()` and `call()` agree numerically under any dtype policy.
dt = tf.einsum(
"br,rd->bd", tf.cast(dt, tf.float32), tf.cast(self.dt_proj.kernel, tf.float32)
)
dt = tf.nn.softplus(dt + tf.cast(self.dt_proj.bias, tf.float32))
dA = tf.exp(tf.einsum("bd,dn->bdn", dt, A))
dB = tf.einsum("bd,bn->bdn", dt, tf.cast(B, tf.float32))
x32 = tf.cast(x, tf.float32)
new_ssm_state = tf.cast(ssm_state, tf.float32) * dA + tf.expand_dims(x32, axis=-1) * dB
y = tf.einsum("bdn,bn->bd", new_ssm_state, tf.cast(C, tf.float32))
y = y + tf.cast(self.D, tf.float32) * x32
y = y * tf.nn.silu(tf.cast(z, tf.float32))
out = self.out_proj(tf.cast(y, compute_dtype))
return tf.expand_dims(out, axis=1), new_states, new_ssm_state