File size: 9,898 Bytes
62d3300 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 |
"""Common Haiku modules."""
from collections.abc import Sequence
import contextlib
import numbers
from typing import TypeAlias
import haiku as hk
import jax
import jax.numpy as jnp
import numpy as np
PRECISION: TypeAlias = (
None
| str
| jax.lax.Precision
| tuple[str, str]
| tuple[jax.lax.Precision, jax.lax.Precision]
)
# Useful for mocking in tests.
DEFAULT_PRECISION = None
# Constant from scipy.stats.truncnorm.std(a=-2, b=2, loc=0., scale=1.)
TRUNCATED_NORMAL_STDDEV_FACTOR = np.asarray(
0.87962566103423978, dtype=np.float32
)
class LayerNorm(hk.LayerNorm):
"""LayerNorm module.
Equivalent to hk.LayerNorm but with an extra 'upcast' option that casts
(b)float16 inputs to float32 before computing the layer norm, and then casts
the output back to the input type.
The learnable parameter shapes are also different from Haiku: they are always
vectors rather than possibly higher-rank tensors. This makes it easier
to change the layout whilst keep the model weight-compatible.
"""
def __init__(
self,
*,
axis: int = -1,
create_scale: bool = True,
create_offset: bool = True,
eps: float = 1e-5,
scale_init: hk.initializers.Initializer | None = None,
offset_init: hk.initializers.Initializer | None = None,
use_fast_variance: bool = True,
name: str,
param_axis: int | None = None,
upcast: bool = True,
):
super().__init__(
axis=axis,
create_scale=False,
create_offset=False,
eps=eps,
scale_init=None,
offset_init=None,
use_fast_variance=use_fast_variance,
name=name,
param_axis=param_axis,
)
self.upcast = upcast
self._temp_create_scale = create_scale
self._temp_create_offset = create_offset
def __call__(self, x: jnp.ndarray) -> jnp.ndarray:
dtype = x.dtype
is_16bit = x.dtype in [jnp.bfloat16, jnp.float16]
if self.upcast and is_16bit:
x = x.astype(jnp.float32)
param_axis = self.param_axis[0] if self.param_axis else -1
param_shape = (x.shape[param_axis],)
param_broadcast_shape = [1] * x.ndim
param_broadcast_shape[param_axis] = x.shape[param_axis]
scale = None
offset = None
if self._temp_create_scale:
scale = hk.get_parameter(
'scale', param_shape, x.dtype, init=self.scale_init
)
scale = scale.reshape(param_broadcast_shape)
if self._temp_create_offset:
offset = hk.get_parameter(
'offset', param_shape, x.dtype, init=self.offset_init
)
offset = offset.reshape(param_broadcast_shape)
out = super().__call__(x, scale=scale, offset=offset)
if self.upcast and is_16bit:
out = out.astype(dtype)
return out
def haiku_linear_get_params(
inputs: jax.Array | jax.ShapeDtypeStruct,
*,
num_output: int | Sequence[int],
use_bias: bool = False,
num_input_dims: int = 1,
initializer: str = 'linear',
bias_init: float = 0.0,
transpose_weights: bool = False,
name: str | None = None,
) -> tuple[jax.Array, jax.Array | None]:
"""Get parameters for linear layer.
Parameters will be at least float32 or higher precision.
Arguments:
inputs: The input to the Linear layer. Can be either a JAX array or a
jax.ShapeDtypeStruct.
num_output: The number of output channels. Can be an integer or a sequence
of integers.
use_bias: Whether to create a bias array.
num_input_dims: The number of dimensions to consider as channel dims in the
input.
initializer: The name of the weight initializer to use.
bias_init: A float used to initialize the bias.
transpose_weights: If True, will create a transposed version of the weights.
name: The Haiku namespace to use for the weight and bias.
Returns:
A tuple[weight, bias] if use_bias otherwise tuple[weight, None].
"""
if isinstance(num_output, numbers.Integral):
output_shape = (num_output,)
else:
output_shape = tuple(num_output)
if num_input_dims > 0:
in_shape = inputs.shape[-num_input_dims:]
elif num_input_dims == 0:
in_shape = ()
else:
raise ValueError('num_input_dims must be >= 0.')
weight_init = _get_initializer_scale(initializer, in_shape)
with hk.name_scope(name) if name else contextlib.nullcontext():
if transpose_weights:
weight_shape = output_shape + in_shape
weights = hk.get_parameter(
'weights', shape=weight_shape, dtype=inputs.dtype, init=weight_init
)
else:
weight_shape = in_shape + output_shape
weights = hk.get_parameter(
name='weights',
shape=weight_shape,
dtype=inputs.dtype,
init=weight_init,
)
bias = None
if use_bias:
bias = hk.get_parameter(
name='bias',
shape=output_shape,
dtype=inputs.dtype,
init=hk.initializers.Constant(bias_init),
)
return weights, bias
class Linear(hk.Module):
"""Custom Linear Module.
This differs from the standard Linear in a few ways:
* It supports inputs of arbitrary rank
* It allows to use ntk parametrization
* Initializers are specified by strings
* It allows to explicitly specify which dimension of the input will map to
the tpu sublane/lane dimensions.
"""
def __init__(
self,
num_output: int | Sequence[int],
*,
initializer: str = 'linear',
num_input_dims: int = 1,
use_bias: bool = False,
bias_init: float = 0.0,
precision: PRECISION = None,
fast_scalar_mode: bool = True,
transpose_weights: bool = False,
name: str,
):
"""Constructs Linear Module.
Args:
num_output: number of output channels. Can be tuple when outputting
multiple dimensions.
initializer: What initializer to use, should be one of {'linear', 'relu',
'zeros'}.
num_input_dims: Number of dimensions from the end to project.
use_bias: Whether to include trainable bias (False by default).
bias_init: Value used to initialize bias.
precision: What precision to use for matrix multiplication, defaults to
None.
fast_scalar_mode: Whether to use optimized path for num_input_dims = 0.
transpose_weights: decides whether weights have shape [input, output] or
[output, input], True means [output, input], this is helpful to avoid
padding on the tensors holding the weights.
name: name of module, used for name scopes.
"""
super().__init__(name=name)
if isinstance(num_output, numbers.Integral):
self.output_shape = (num_output,)
else:
self.output_shape = tuple(num_output)
self.initializer = initializer
self.use_bias = use_bias
self.bias_init = bias_init
self.num_input_dims = num_input_dims
self.num_output_dims = len(self.output_shape)
self.precision = precision if precision is not None else DEFAULT_PRECISION
self.fast_scalar_mode = fast_scalar_mode
self.transpose_weights = transpose_weights
def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
"""Connects Module.
Args:
inputs: Tensor of shape [..., num_channel]
Returns:
output of shape [..., num_output]
"""
num_input_dims = self.num_input_dims
# Adds specialized path for scalar inputs in Linear layer,
# this means the linear Layer does not use the matmul units on the tpu,
# which is more efficient and gives compiler more flexibility over layout.
if num_input_dims == 0 and self.fast_scalar_mode:
weight_shape = self.output_shape
if self.initializer == 'zeros':
w_init = hk.initializers.Constant(0.0)
else:
distribution_stddev = jnp.array(1 / TRUNCATED_NORMAL_STDDEV_FACTOR)
w_init = hk.initializers.TruncatedNormal(
mean=0.0, stddev=distribution_stddev
)
weights = hk.get_parameter('weights', weight_shape, inputs.dtype, w_init)
inputs = jnp.expand_dims(
inputs, tuple(range(-1, -self.num_output_dims - 1, -1))
)
output = inputs * weights
else:
if self.num_input_dims > 0:
in_shape = inputs.shape[-self.num_input_dims :]
else:
in_shape = ()
weight_init = _get_initializer_scale(self.initializer, in_shape)
in_letters = 'abcde'[: self.num_input_dims]
out_letters = 'hijkl'[: self.num_output_dims]
if self.transpose_weights:
weight_shape = self.output_shape + in_shape
weights = hk.get_parameter(
'weights', weight_shape, inputs.dtype, weight_init
)
equation = (
f'...{in_letters}, {out_letters}{in_letters}->...{out_letters}'
)
else:
weight_shape = in_shape + self.output_shape
weights = hk.get_parameter(
'weights', weight_shape, inputs.dtype, weight_init
)
equation = (
f'...{in_letters}, {in_letters}{out_letters}->...{out_letters}'
)
output = jnp.einsum(equation, inputs, weights, precision=self.precision)
if self.use_bias:
bias = hk.get_parameter(
'bias',
self.output_shape,
inputs.dtype,
hk.initializers.Constant(self.bias_init),
)
output += bias
return output
def _get_initializer_scale(initializer_name, input_shape):
"""Get initializer for weights."""
if initializer_name == 'zeros':
w_init = hk.initializers.Constant(0.0)
else:
# fan-in scaling
noise_scale = 1.0
for channel_dim in input_shape:
noise_scale /= channel_dim
if initializer_name == 'relu':
noise_scale *= 2
stddev = np.sqrt(noise_scale)
# Adjust stddev for truncation.
stddev = stddev / TRUNCATED_NORMAL_STDDEV_FACTOR
w_init = hk.initializers.TruncatedNormal(mean=0.0, stddev=stddev)
return w_init
|