Variance and InverseVariance uncertainty propagation escapes to numpy on jax (only the StdDev wrapper is array-API-aware)

[mwcraig]

Behavior

Under the jax backend, arithmetic on a wrapped CCDData whose uncertainty is VarianceUncertainty or InverseVariance returns an uncertainty whose .array is a numpy array, while the data stays jax. StdDevUncertainty correctly stays jax. First public call site is the gain path added in #958: cosmicray_lacosmic(ccd, gain=..., gain_apply=True) with a Variance or InverseVariance uncertainty on jax gives result.uncertainty.array as numpy (verified 2026-07-29 against the #958 branch; the values are numerically correct — this is an escape, not a wrong result). On CuPy the same path would mean an implicit device transfer or failure rather than a silent conversion.

Root cause

In _ccddata_wrapper_for_array_api.py, only _StdDevUncertaintyWrapper overrides the four _propagate_* methods to pass namespace-aware hooks into astropy's generic propagation (to_variance=xp.square, from_variance=xp.sqrt):

ccdproc/ccdproc/_ccddata_wrapper_for_array_api.py

Lines 266 to 315 in 736555d

	class _StdDevUncertaintyWrapper(_CupyOperationNamesMixin, StdDevUncertainty):
	"""
	Override operation propagate methods to make sure they use the array API.
	Override overall propagate method to allow cupy_-prefixed operation names.
	"""
	def _propagate_add(self, other_uncert, result_data, correlation):
	xp = array_api_compat.array_namespace(self.array, other_uncert.array)
	return super()._propagate_add_sub(
	other_uncert,
	result_data,
	correlation,
	subtract=False,
	to_variance=xp.square,
	from_variance=xp.sqrt,
	)
	def _propagate_subtract(self, other_uncert, result_data, correlation):
	xp = array_api_compat.array_namespace(self.array, other_uncert.array)
	return super()._propagate_add_sub(
	other_uncert,
	result_data,
	correlation,
	subtract=True,
	to_variance=xp.square,
	from_variance=xp.sqrt,
	)
	def _propagate_multiply(self, other_uncert, result_data, correlation):
	xp = array_api_compat.array_namespace(self.array, other_uncert.array)
	return super()._propagate_multiply_divide(
	other_uncert,
	result_data,
	correlation,
	divide=False,
	to_variance=xp.square,
	from_variance=xp.sqrt,
	)
	def _propagate_divide(self, other_uncert, result_data, correlation):
	xp = array_api_compat.array_namespace(self.array, other_uncert.array)
	return super()._propagate_multiply_divide(
	other_uncert,
	result_data,
	correlation,
	divide=True,
	to_variance=xp.square,
	from_variance=xp.sqrt,
	)

_VarianceUncertaintyWrapper and _InverseVarianceWrapper are operation-name shims only — their _propagate_* methods fall through to astropy's stock numpy implementations, which coerce jax arrays:

ccdproc/ccdproc/_ccddata_wrapper_for_array_api.py

Lines 318 to 323 in 736555d

	class _VarianceUncertaintyWrapper(_CupyOperationNamesMixin, VarianceUncertainty):
	"""This subclass is needed to allow CuPy operation names"""
	class _InverseVarianceWrapper(_CupyOperationNamesMixin, InverseVariance):
	"""This subclass is needed to allow CuPy operation names"""

Current tracking

Only two backend_xfail("jax") markers on test_cosmicray_gain_correct_uncertainty_namespace (added in #958). The xfails are not strict, so a fix would silently start xpassing rather than flag itself — this issue is the durable record.

Suggested fix

Mirror the StdDev pattern in the other two wrappers: override the four _propagate_* methods to pass xp-aware to_variance/from_variance callables into _propagate_add_sub/_propagate_multiply_divide (identity for Variance, reciprocal for InverseVariance). Then flip the two jax xfail markers to assertions (or make them strict=True first).

Follow-up to #909. This issue was written by Claude (via Claude Code); Matt reviewed and approved its content before filing.