Bases: QNumerics, PyTreeNode
No quantization, use a native type such as bf16.
Source code in src/fjformer/bits/no_numerics.py
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45 | class NoNumerics(numerics.QNumerics, flax.struct.PyTreeNode):
"""No quantization, use a native type such as bf16."""
# TODO(lew): This is a workaround. We should separate Stochastic Rounding.
# noise_fn has no effect in NoNumerics.
noise_fn: Optional[stochastic_rounding.NoiseFn] = None
dtype: Optional[Any] = None
# TODO(lew): This is a hack. We treat check isinstance(NoNumerics) and treat
# it in a special way right now. These functions are never called
def get_dtype(self):
return None
def fwd(self, x, context):
pass
def abs_val_mapped_to(self):
pass
def vjp_fwd(self, x, context):
pass
def vjp_bwd(self, res, grad):
pass
|