Parametrise type of static params to exclude from dynamic param creat…

…ion in sources.

This allows us to extend the exclusion logic more simply for child sources.

PiperOrigin-RevId: 711463448

torax/sources/runtime_params.py

@@ -99,6 +99,25 @@ def build_static_params(self) -> StaticRuntimeParams:
     )
 
 
+@chex.dataclass(frozen=True)
+class DynamicRuntimeParams:
+  """Dynamic params for a single TORAX source.
+
+  These params can be changed without triggering a recompile. TORAX sources are
+  stateless, so these params are their inputs to determine their output
+  profiles.
+  """
+  prescribed_values: array_typing.ArrayFloat
+
+
+@chex.dataclass(frozen=True)
+class StaticRuntimeParams:
+  """Static params for the sources."""
+
+  mode: int
+  is_explicit: bool
+
+
 @chex.dataclass
 class RuntimeParamsProvider(
     base.RuntimeParametersProvider['DynamicRuntimeParams']
@@ -111,12 +130,15 @@ class RuntimeParamsProvider(
   def get_dynamic_params_kwargs(
       self,
       t: chex.Numeric,
+      static_runtime_params_type: type[
+          StaticRuntimeParams
+      ] = StaticRuntimeParams,
   ) -> dict[str, Any]:
     dynamic_params_kwargs = super(
         RuntimeParamsProvider, self
     ).get_dynamic_params_kwargs(t)
     # Remove any fields from runtime params that are included in static params.
-    for field in dataclasses.fields(StaticRuntimeParams):
+    for field in dataclasses.fields(static_runtime_params_type):
       del dynamic_params_kwargs[field.name]
     return dynamic_params_kwargs
 
@@ -125,22 +147,3 @@ def build_dynamic_params(
       t: chex.Numeric,
   ) -> DynamicRuntimeParams:
     return DynamicRuntimeParams(**self.get_dynamic_params_kwargs(t))
-
-
-@chex.dataclass(frozen=True)
-class DynamicRuntimeParams:
-  """Dynamic params for a single TORAX source.
-
-  These params can be changed without triggering a recompile. TORAX sources are
-  stateless, so these params are their inputs to determine their output
-  profiles.
-  """
-  prescribed_values: array_typing.ArrayFloat
-
-
-@chex.dataclass(frozen=True)
-class StaticRuntimeParams:
-  """Static params for the sources."""
-
-  mode: int
-  is_explicit: bool