Problem
apply() and compress() take a dtype argument that defaults to
np.float64 and is documented as "The data type for the computation". Neither
statement matches the behaviour:
dtype is used for exactly one thing — the dtype of the random sketch
omega. It has no effect on the dtype of the result, which is determined by
NumPy type promotion against the input tensors.
- Because the default is
np.float64, a single-precision input is silently
promoted to double precision:
A = qtn.MPO_rand(8, bond_dim=32, phys_dim=2, dtype=np.complex64)
compress(A, chi_out=32).arrays[2].dtype
# complex128 <-- 2x the memory and bandwidth the caller asked for
compress(A, chi_out=32, dtype=np.float32).arrays[2].dtype
# complex64 <-- only correct if the caller knows to pass this
This is precisely the wrong default for the GPU path, where single precision is
the main reason to use a GPU at all. Users silently lose half their throughput
and double their memory unless they happen to discover the dtype argument.
Note also that dtype accepts a type, so np.complex128 is a legal value;
that just makes omega a complex array with a zero imaginary part, doubling
the sketch cost for no accuracy benefit. A real Gaussian sketch is sufficient
for a complex operator, as confirmed by measurement: the compression error is
identical either way (2.98e-08 in both cases).
Proposed fix
Derive the sketch dtype from the input instead of from a fixed default:
- Compute the promoted input dtype (e.g.
np.result_type(*mpo.arrays)), then
draw omega in the corresponding real floating type
(np.finfo(input_dtype).dtype), so a complex64 input gets a float32
sketch and stays in complex64 end to end.
- Change the signature to
dtype: DTypeLike | None = None, where None means
"follow the input". Keep the argument as an escape hatch for callers who
deliberately want to sketch in a different precision.
- Fix the docstring to say what the parameter actually controls.
Acceptance criteria
Problem
apply()andcompress()take adtypeargument that defaults tonp.float64and is documented as "The data type for the computation". Neitherstatement matches the behaviour:
dtypeis used for exactly one thing — the dtype of the random sketchomega. It has no effect on the dtype of the result, which is determined byNumPy type promotion against the input tensors.
np.float64, a single-precision input is silentlypromoted to double precision:
This is precisely the wrong default for the GPU path, where single precision is
the main reason to use a GPU at all. Users silently lose half their throughput
and double their memory unless they happen to discover the
dtypeargument.Note also that
dtypeaccepts atype, sonp.complex128is a legal value;that just makes
omegaa complex array with a zero imaginary part, doublingthe sketch cost for no accuracy benefit. A real Gaussian sketch is sufficient
for a complex operator, as confirmed by measurement: the compression error is
identical either way (
2.98e-08in both cases).Proposed fix
Derive the sketch dtype from the input instead of from a fixed default:
np.result_type(*mpo.arrays)), thendraw
omegain the corresponding real floating type(
np.finfo(input_dtype).dtype), so acomplex64input gets afloat32sketch and stays in
complex64end to end.dtype: DTypeLike | None = None, whereNonemeans"follow the input". Keep the argument as an escape hatch for callers who
deliberately want to sketch in a different precision.
Acceptance criteria
complex64in ⇒complex64out;float32in ⇒float32out, with no argument passed.complex128/float64behaviour is unchanged.dtype.