braintools 0.2.0__tar.gz → 0.3.0__tar.gz

{braintools-0.2.0 → braintools-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintools
-Version: 0.2.0
+Version: 0.3.0
 Summary: Modeling tools for brain simulation.
 Author-email: BrainTools Developers <chao.brain@qq.com>
 License: 

{braintools-0.2.0 → braintools-0.3.0}/braintools/_version.py

@@ -13,5 +13,5 @@
 # limitations under the License.
 # ==============================================================================
 
-__version__ = "0.2.0"
+__version__ = "0.3.0"
 __version_info__ = tuple(map(int, __version__.split(".")))

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/__init__.py

@@ -35,7 +35,7 @@ Building custom tasks from phases:
 >>> from braintools.cogtask import (
 ...     Task, Context, concat,
 ...     Fixation, Stimulus, Delay, Response,
-...     Feature, circular, one_hot
+...     Feature, circular, one_hot, label
 ... )
 >>> import brainunit as u
 >>>
@@ -44,13 +44,14 @@ Building custom tasks from phases:
 >>> stim = Feature(8, 'stimulus')
 >>> choice = Feature(2, 'choice')
 >>>
->>> # Build task
+>>> # Build task. A string label spec must be wrapped in ``label(...)`` so it
+>>> # is read as a context key; a bare string is treated as a literal and fails.
 >>> task = Task(
 ...     phases=concat([
 ...         Fixation(100 * u.ms, inputs={'fixation': 1.0}),
 ...         Stimulus(500 * u.ms, inputs={'stimulus': circular('direction')}),
 ...         Delay(500 * u.ms, inputs={'fixation': 1.0}),
-...         Response(100 * u.ms, outputs={'label': 'ground_truth'})
+...         Response(100 * u.ms, outputs={'label': label('ground_truth')})
 ...     ]),
 ...     input_features=fix + stim,
 ...     output_features=fix + choice,
@@ -153,7 +154,7 @@ from .phase import (
     execute_phase_packed,
     phase_tree_is_variable,
 )
-from .task import Task
+from .task import Task, create_task
 # Pre-built tasks
 from .tasks import (
     # Decision Making
@@ -215,6 +216,7 @@ __all__ = [
     'Switch',
     'While',
     'Task',
+    'create_task',
     # Features
     'Feature',
     'FeatureSet',

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/context.py

@@ -36,19 +36,25 @@ class Context:
 
     Examples
     --------
-    >>> ctx = Context(seed=42)
-    >>> ctx['ground_truth'] = 1
-    >>> ctx['stimulus_direction'] = 0.5 * np.pi
-    >>> print(ctx['ground_truth'])
-    1
+    .. code-block:: python
 
-    >>> # In a phase:
-    >>> ctx.inputs[ctx.phase_start:ctx.phase_end, :] = stimulus_encoding
+        >>> import jax
+        >>> import jax.numpy as jnp
+        >>> ctx = Context(key=jax.random.PRNGKey(42))
+        >>> ctx['ground_truth'] = 1
+        >>> ctx['stimulus_direction'] = 0.5 * jnp.pi
+        >>> print(ctx['ground_truth'])
+        1
+
+        >>> # In a phase, write into the buffer with a functional update
+        >>> # (JAX arrays are immutable; ``.at[...].set(...)`` returns a copy):
+        >>> # ctx.inputs = ctx.inputs.at[ctx.phase_start:ctx.phase_end, :].set(stimulus_encoding)
 
     Parameters
     ----------
     key : jax.Array, optional
-        JAX PRNGKey for random number generation. If None, creates one with seed.
+        JAX PRNGKey for random number generation. If None, draws fresh
+        randomness from brainstate's default RNG.
 
     Note
     ----

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/feature.py

@@ -300,7 +300,7 @@ class Feature(_FeatureBase):
         """
         Create multiple copies of this feature with indexed names.
 
-        Example: Feature(1, 30.*u.Hz, 'choice') * 3
+        Example: Feature(1, 'choice') * 3
         Creates: choice_0, choice_1, choice_2
         """
         if not isinstance(count, int) or count < 1:

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/phase.py

@@ -457,25 +457,33 @@ class Parallel(Phase):
         pass
 
     def execute(self, ctx: Context) -> None:
-        """Encode each child within its own [phase_start, phase_start+dur)."""
+        """Encode each child within its own ``[parent_start, parent_start+dur)``.
+
+        Every child — leaf *or* compound — is dispatched through
+        ``execute_phase`` so that nested ``Sequence``/``Repeat``/conditional
+        children execute correctly. Routing compound children through
+        ``encode_inputs``/``encode_outputs`` directly would be a no-op (those
+        classes drive their sub-phases via ``execute``), silently dropping the
+        branch. This mirrors ``execute_packed`` and keeps the two paths
+        consistent. All children start at ``parent_start``; only the first child
+        contributes to the output buffer by convention.
+        """
         parent_start = ctx.phase_start
         parent_end = ctx.phase_end
         parent_name = ctx.current_phase
 
         for i, child in enumerate(self.phases):
-            child_dur = child.get_duration(ctx)
-            ctx.phase_start = parent_start
-            ctx.phase_end = parent_start + child_dur
-            ctx.current_phase = child.name
-            child.on_enter(ctx)
-            child.encode_inputs(ctx)
-            # Only the first child contributes outputs by default.
+            ctx.current_step = parent_start
             if i == 0:
-                child.encode_outputs(ctx)
-            child.on_exit(ctx)
-            ctx.phase_history.append((child.name, ctx.phase_start, ctx.phase_end))
+                execute_phase(child, ctx)
+            else:
+                # Non-first children must not contribute to the output buffer.
+                # Stash, encode (inputs + any state), then restore outputs.
+                saved_out = ctx.outputs
+                execute_phase(child, ctx)
+                ctx.outputs = saved_out
 
-        # Restore parent scope.
+        # Restore parent scope and advance past the (longest) child.
         ctx.phase_start = parent_start
         ctx.phase_end = parent_end
         ctx.current_phase = parent_name

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/phase_test.py

@@ -333,3 +333,63 @@ def test_describe_lists_inputs_outputs_and_noise():
     s = p.describe()
     assert "Inputs" in s and "Outputs" in s and "Noise" in s
     assert "fixation" in s
+
+
+def test_parallel_executes_compound_child_eager():
+    # Regression: a compound child (Sequence) of a Parallel must execute its
+    # sub-phases. Previously Parallel.execute called encode_inputs directly,
+    # which is a no-op for compounds, so the whole branch was silently dropped.
+    m = ct.Feature(1, 'm')
+    n = ct.Feature(1, 'n')
+    fix = ct.Feature(1, 'fixation')
+    a = Stimulus(5 * u.ms, name='A', inputs={'m': 1.0}, outputs={'label': 0})
+    b = Stimulus(5 * u.ms, name='B', inputs={'m': 1.0}, outputs={'label': 0})
+    seq_ab = a >> b                      # compound child
+    c = Stimulus(10 * u.ms, name='C', inputs={'n': 1.0}, outputs={'label': 0})
+    par = seq_ab | c
+    task = ct.Task(
+        phases=par,
+        input_features=m + n + fix,
+        output_features=fix + ct.Feature(1, 'out'),
+        seed=0,
+    )
+    X, _, _ = task.sample_trial(0)
+    X = np.asarray(X)
+    assert X.shape[0] == 10
+    # A and B each ran for 5 ticks writing 1.0 into channel 'm' -> sum == 10.
+    np.testing.assert_allclose(X[:, 0].sum(), 10.0, atol=1e-6)
+    # C ran for all 10 ticks writing 1.0 into channel 'n'.
+    np.testing.assert_allclose(X[:, 1], 1.0, atol=1e-6)
+
+
+def test_parallel_eager_and_packed_agree_on_compound_child():
+    # The fixed (eager) and variable (packed) execution paths must produce the
+    # same input channels for a Parallel whose first child is a compound, when
+    # all durations are fully realised.
+    m = ct.Feature(1, 'm')
+    n = ct.Feature(1, 'n')
+    fix = ct.Feature(1, 'fixation')
+
+    def build(extra_phase):
+        a = Stimulus(5 * u.ms, name='A', inputs={'m': 1.0}, outputs={'label': 0})
+        b = Stimulus(5 * u.ms, name='B', inputs={'m': 1.0}, outputs={'label': 0})
+        c = Stimulus(10 * u.ms, name='C', inputs={'n': 1.0}, outputs={'label': 0})
+        phases = (a >> b) | c
+        if extra_phase is not None:
+            phases = phases >> extra_phase
+        return ct.Task(phases=phases, input_features=m + n + fix,
+                       output_features=fix + ct.Feature(1, 'out'), seed=0)
+
+    # Eager path.
+    X_eager, _, _ = build(None).sample_trial(0)
+    # Packed path: append a VariableDuration whose realised length is fixed via
+    # trial_init, forcing packed mode but the same total realised content.
+    vd = ct.VariableDuration(min_duration=2 * u.ms, max_duration=4 * u.ms,
+                             ctx_key='d', inputs={'n': 1.0}, outputs={'label': 0},
+                             name='vd')
+    packed_task = build(vd)
+    packed_task._trial_init_func = lambda ctx: ctx.update(d=4.0)
+    X_packed, _, info = packed_task.sample_trial(0)
+    # The Parallel block occupies the first 10 ticks in both; compare them.
+    np.testing.assert_allclose(np.asarray(X_eager)[:10, 0],
+                               np.asarray(X_packed)[:10, 0], atol=1e-6)

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/task.py

@@ -86,17 +86,22 @@ class Task:
     --------
     Instance-based (traditional):
 
+    >>> import brainunit as u
+    >>> from braintools.cogtask import Feature, Fixation, Stimulus, Response, circular, label
+    >>> fix = Feature(1, 'fixation')
+    >>> stim = Feature(8, 'stimulus')
+    >>> choice = Feature(2, 'choice')
     >>> task = Task(
     ...     phases=(
-    ...         Fixation(100 * u.ms)
-    ...         >> Stimulus(2000 * u.ms, feature=stim, encoder=circular_encoder())
-    ...         >> Response(100 * u.ms, output_feature=choice)
+    ...         Fixation(100 * u.ms, inputs={'fixation': 1.0})
+    ...         >> Stimulus(2000 * u.ms, inputs={'stimulus': circular('direction')})
+    ...         >> Response(100 * u.ms, outputs={'label': label('ground_truth')})
     ...     ),
     ...     input_features=fix + stim,
     ...     output_features=fix + choice,
     ...     trial_init=lambda ctx: ctx.update(
     ...         ground_truth=ctx.rng.choice(2),
-    ...         direction=ctx.rng.uniform(0, 2*np.pi)
+    ...         direction=ctx.rng.uniform(0, 2 * 3.14159)
     ...     )
     ... )
 
@@ -107,12 +112,13 @@ class Task:
     ...     num_stimuli = 8
     ...
     ...     def define_features(self):
-    ...         fix = Feature(1, 40*u.Hz, 'fixation')
-    ...         stim = Feature(self.num_stimuli, 40*u.Hz, 'stimulus')
-    ...         return fix + stim, fix + Feature(2, 40*u.Hz, 'response')
+    ...         fix = Feature(1, 'fixation')
+    ...         stim = Feature(self.num_stimuli, 'stimulus')
+    ...         return fix + stim, fix + Feature(2, 'response')
     ...
     ...     def define_phases(self):
-    ...         return FixationPhase(self.t_fixation) >> ResponsePhase()
+    ...         return (Fixation(self.t_fixation, inputs={'fixation': 1.0})
+    ...                 >> Response(100 * u.ms, outputs={'label': label('ground_truth')}))
     ...
     ...     def trial_init(self, ctx):
     ...         ctx['ground_truth'] = ctx.rng.choice(2)
@@ -631,11 +637,20 @@ def create_task(
 
     Examples
     --------
+    >>> import brainunit as u
+    >>> from braintools.cogtask import Feature, Fixation, Stimulus, Response, circular, label
     >>> task = create_task(
-    ...     phases=Fixation(100*u.ms) >> Stimulus(500*u.ms) >> Response(100*u.ms),
-    ...     input_features=Feature('fix', 1) + Feature('stim', 2),
-    ...     output_features=Feature('fix', 1) + Feature('choice', 2),
-    ...     num_trial=1000
+    ...     phases=(
+    ...         Fixation(100 * u.ms, inputs={'fixation': 1.0})
+    ...         >> Stimulus(500 * u.ms, inputs={'stimulus': circular('direction')})
+    ...         >> Response(100 * u.ms, outputs={'label': label('ground_truth')})
+    ...     ),
+    ...     input_features=Feature(1, 'fixation') + Feature(2, 'stimulus'),
+    ...     output_features=Feature(1, 'fixation') + Feature(2, 'choice'),
+    ...     trial_init=lambda ctx: ctx.update(
+    ...         ground_truth=ctx.rng.choice(2),
+    ...         direction=ctx.rng.uniform(0, 2 * 3.14159)
+    ...     ),
     ... )
     """
     return Task(phases, input_features, output_features, trial_init, name,

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/task_test.py

@@ -387,3 +387,36 @@ def test_pulse_decision_making_runs_through_repeat():
         gt = int(info['trial_state']['ground_truth'])
         last = int(np.asarray(Y)[-1])
         assert last == gt + 1
+
+
+# ---------------------------------------------------------------------------
+# create_task factory: must be exported from the package and produce a
+# sampleable task. Also exercises the (previously broken) docstring pattern
+# of a string label spec wrapped in label(...).
+# ---------------------------------------------------------------------------
+
+def test_create_task_is_exported_from_package():
+    assert hasattr(ct, 'create_task')
+    assert 'create_task' in ct.__all__
+
+
+def test_create_task_builds_sampleable_task_with_string_label():
+    task = ct.create_task(
+        phases=(
+            ct.Fixation(10 * u.ms, inputs={'fixation': 1.0})
+            >> ct.Stimulus(20 * u.ms, inputs={'stimulus': ct.circular('direction')})
+            >> ct.Response(10 * u.ms, outputs={'label': ct.label('ground_truth')})
+        ),
+        input_features=ct.Feature(1, 'fixation') + ct.Feature(2, 'stimulus'),
+        output_features=ct.Feature(1, 'fixation') + ct.Feature(2, 'choice'),
+        trial_init=lambda ctx: ctx.update(
+            ground_truth=ctx.rng.choice(2),
+            direction=ctx.rng.uniform(0, 2 * 3.14159),
+        ),
+        seed=0,
+    )
+    X, Y, info = task.sample_trial(0)
+    assert X.shape[0] == Y.shape[0] == 40
+    gt = int(info['trial_state']['ground_truth'])
+    # Response phase is the last 10 ticks; label == ground_truth there.
+    assert (np.asarray(Y)[-10:] == gt).all()

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/tasks/working_memory.py

@@ -62,6 +62,13 @@ def make_encoder(mode: str, key: str, *, feature_per_direction: int = 1, kappa=2
     """
     mode = mode.lower()
 
+    # ``scalar`` broadcasts a single value across every feature dimension, so it
+    # is agnostic to ``feature_per_direction`` — handle it up front rather than
+    # only in the K == 1 branch (otherwise scalar + K>1 fell through to the
+    # population branch and raised a misleading "Unknown mode=scalar").
+    if mode == "scalar":
+        return scalar(key)
+
     if feature_per_direction == 1:
         if mode == "von_mises":
             return von_mises(key, kappa=kappa, base_value=base_value, as_index=True, num_dirs=num_dirs)
@@ -69,8 +76,6 @@ def make_encoder(mode: str, key: str, *, feature_per_direction: int = 1, kappa=2
             return one_hot(key)
         elif mode == "circular":
             return circular(key, base_value=base_value, as_index=True, num_dirs=num_dirs)
-        elif mode == "scalar":
-            return scalar(key)
         else:
             raise ValueError(f"Unknown mode={mode}")
 

{braintools-0.2.0 → braintools-0.3.0}/braintools/cogtask/tasks/working_memory_test.py

@@ -67,6 +67,20 @@ def test_make_encoder_population_dim_not_divisible_raises():
         enc(ctx, f)
 
 
+@pytest.mark.parametrize("k", [1, 3, 5])
+def test_make_encoder_scalar_is_feature_per_direction_agnostic(k):
+    # Regression: scalar was only handled in the K==1 branch, so
+    # feature_per_direction>1 raised "Unknown mode=scalar". scalar broadcasts a
+    # single value across all dims and must work for any K.
+    enc = make_encoder("scalar", "sample_value", feature_per_direction=k)
+    f = ct.Feature(6, 'stimulus')
+    ctx = ct.Context()
+    ctx['sample_value'] = 0.7
+    out = np.asarray(enc(ctx, f))
+    assert out.shape == (6,)
+    np.testing.assert_allclose(out, 0.7, atol=1e-6)
+
+
 def test_build_cues_returns_jax_arrays_with_right_shapes():
     non_resp, resp = build_cues(3)
     assert non_resp.shape == (3,)

{braintools-0.2.0 → braintools-0.3.0}/braintools/init/_init_base.py

@@ -87,9 +87,13 @@ class Initialization(ABC):
         >>>
         >>> delay_init = Uniform(1.0 * u.ms, 3.0 * u.ms).clip(0.5 * u.ms, 5.0 * u.ms)
         >>>
+        >>> # Each transform must be parenthesized: a bare ``lambda`` body would
+        >>> # greedily capture the rest of the expression. Bounds must share the
+        >>> # value's unit (``0.0 * u.nS``, not ``0``).
         >>> combined = (Normal(1.0 * u.nS, 0.2 * u.nS) |
-        ...             lambda x: x.clip(0, 2 * u.nS) |
-        ...             lambda x: x * 0.5)
+        ...             (lambda x: u.math.maximum(x, 0.0 * u.nS)) |
+        ...             (lambda x: x * 0.5))
+        >>> values = combined(1000, rng=np.random.default_rng(0))
     """
     __module__ = 'braintools.init'
 
@@ -344,9 +348,17 @@ def param(
         return init
 
     # Check if the parameter is a callable function
+    #
+    # For a callable, the batch axis is incorporated by prepending it to ``sizes``
+    # and letting the initializer fill the full batched shape (so a random
+    # initializer draws independently per batch element). The shared trailing
+    # batch-expansion below must therefore NOT run again for this path, otherwise
+    # the batch axis is applied twice (``(batch, batch, *sizes)``).
+    batch_applied_by_callable = False
     if callable(init):
         if batch_size is not None:
             sizes = (batch_size,) + sizes
+            batch_applied_by_callable = True
         init = init(sizes, **param_kwargs)
     else:
         # Check if the parameter is a scalar value
@@ -357,16 +369,20 @@ def param(
         else:
             raise TypeError(f'Unknown parameter type: {type(init)}')
 
+    # Resolve the underlying value first: a ``State`` is a mutable container and
+    # ``u.math.shape(state)`` reports ``()`` (it does not see into the wrapped
+    # value), which would silently bypass the shape check below.
+    param_value = init.value if isinstance(init, State) else init
+
     # Check if the shape of the parameter matches the given size
-    if not _are_broadcastable_shapes(u.math.shape(init), sizes):
+    if not _are_broadcastable_shapes(u.math.shape(param_value), sizes):
         raise ValueError(
-            f'The shape of the parameter {u.math.shape(init)} '
+            f'The shape of the parameter {u.math.shape(param_value)} '
             f'does not match with the given size {sizes}'
         )
 
     # Expand the parameter to match the given batch size
-    param_value = init.value if isinstance(init, State) else init
-    if batch_size is not None:
+    if batch_size is not None and not batch_applied_by_callable:
         if param_value.ndim <= len(sizes):
             # add a new axis to the params so that it matches the dimensionality of the given shape ``sizes``
             param_value = _expand_params_to_match_sizes(param_value, sizes)
@@ -375,15 +391,19 @@ def param(
                 batch_size,
                 axis=0
             )
-            if isinstance(init, State):
-                init.restore_value(param_value)
-                param_value = init
         else:
             if param_value.shape[0] != batch_size:
                 raise ValueError(
                     f'The batch size of the parameter {param_value.shape[0]} '
                     f'does not match with the given batch size {batch_size}'
                 )
+    # A ``State`` input is a mutable container: update its value and return the
+    # container itself, consistent with the ``Param`` path above and the
+    # documented contract. Doing this only in the batch-expansion branch (the
+    # previous behaviour) made the return type depend on the batch shape.
+    if isinstance(init, State):
+        init.restore_value(param_value)
+        return init
     return param_value
 
 

{braintools-0.2.0 → braintools-0.3.0}/braintools/init/_init_base_test.py

@@ -709,6 +709,50 @@ class TestParamSupport:
             param(p, 5, batch_size=3)
 
 
+class TestParamStateSupport:
+    """param() must return the State (value updated) consistently (bug B2).
+
+    Previously the no-batch and already-batched paths returned the bare
+    underlying value, while only the batch-expansion path returned the State.
+    """
+
+    def test_state_no_batch_returns_state(self):
+        from brainstate import State
+        s = State(np.ones(10) * u.nS)
+        result = param(s, 10)
+        assert isinstance(result, State)
+        assert result is s
+        assert u.math.shape(result.value) == (10,)
+
+    def test_state_batch_expand_returns_state(self):
+        from brainstate import State
+        s = State(np.ones(10) * u.nS)
+        result = param(s, 10, batch_size=4)
+        assert isinstance(result, State)
+        assert result is s
+        assert u.math.shape(result.value) == (4, 10)
+
+    def test_state_already_batched_returns_state(self):
+        from brainstate import State
+        s = State(np.ones((4, 10)) * u.nS)
+        result = param(s, 10, batch_size=4)
+        assert isinstance(result, State)
+        assert result is s
+        assert u.math.shape(result.value) == (4, 10)
+
+    def test_state_batch_mismatch_raises(self):
+        from brainstate import State
+        s = State(np.ones((3, 10)) * u.nS)
+        with pytest.raises(ValueError, match="batch size"):
+            param(s, 10, batch_size=4)
+
+    def test_state_shape_mismatch_raises(self):
+        from brainstate import State
+        s = State(np.ones(10) * u.nS)
+        with pytest.raises(ValueError, match="does not match"):
+            param(s, 5)
+
+
 class TestParamBatching:
     """Test param() batch_size expansion for plain arrays/quantities (bug C2)."""
 
@@ -741,6 +785,59 @@ class TestParamBatching:
             param(np.ones((4, 5)), 5, batch_size=3)
 
 
+class TestParamCallableBatching:
+    """param() batch_size for callable / Initialization initializers.
+
+    Regression for the double-batching bug: a callable initializer with
+    ``batch_size`` must yield exactly ``(batch_size, *sizes)`` -- not
+    ``(batch_size, batch_size, *sizes)``.
+    """
+
+    def test_callable_with_batch_size_1d(self):
+        """A callable gains a single leading batch axis."""
+        result = param(SimpleInit(5.0), 3, batch_size=4)
+        assert u.math.shape(result) == (4, 3)
+        assert np.all(np.asarray(result) == 5.0)
+
+    def test_callable_with_batch_size_2d(self):
+        """A callable with 2-D sizes gains exactly one leading batch axis."""
+        result = param(SimpleInit(5.0), (3, 4), batch_size=2)
+        assert u.math.shape(result) == (2, 3, 4)
+        assert np.all(np.asarray(result) == 5.0)
+
+    def test_callable_quantity_with_batch_size(self):
+        """A callable returning a Quantity keeps its unit and batched shape."""
+        result = param(QuantityInit(2.0 * u.nS), 3, batch_size=4)
+        assert u.math.shape(result) == (4, 3)
+        assert u.get_unit(result) == u.nS
+        assert np.all(result.to_decimal(u.nS) == 2.0)
+
+    def test_callable_batch_size_none_unchanged(self):
+        """batch_size=None leaves the callable result un-batched."""
+        result = param(SimpleInit(5.0), 3, batch_size=None)
+        assert u.math.shape(result) == (3,)
+
+    def test_callable_random_independent_per_batch(self):
+        """A random initializer draws independently for each batch element."""
+        rng = np.random.default_rng(0)
+        result = param(RandomInit(0.0, 1.0), 3, batch_size=4, rng=rng)
+        assert u.math.shape(result) == (4, 3)
+        arr = np.asarray(result)
+        # Independent draws => batch rows must not all be identical.
+        assert not np.allclose(arr[0], arr[1])
+
+    def test_callable_forwards_kwargs_with_batch_size(self):
+        """**param_kwargs still reach the callable when batching."""
+
+        class KwargsInit(Initialization):
+            def __call__(self, size, fill=None, **kwargs):
+                return np.full(size, fill)
+
+        result = param(KwargsInit(), 3, batch_size=4, fill=7.0)
+        assert u.math.shape(result) == (4, 3)
+        assert np.all(np.asarray(result) == 7.0)
+
+
 class TestPipeInitRejectsInitialization:
     """PipeInit cannot pipe values into an Initialization right operand (bug H3)."""
 

{braintools-0.2.0 → braintools-0.3.0}/braintools/init/_init_basic.py

@@ -482,21 +482,27 @@ class TruncatedNormal(Initialization):
         lo = None if self.low is None else u.Quantity(self.low).to(unit).mantissa
         hi = None if self.high is None else u.Quantity(self.high).to(unit).mantissa
 
-        if std == 0:
-            # Degenerate distribution: all mass at the (clamped) mean.
-            samples = jnp.full(size, mean, dtype=float)
-        else:
-            a = -jnp.inf if lo is None else (lo - mean) / std
-            b = jnp.inf if hi is None else (hi - mean) / std
-            # Inverse-CDF sampling: draw u ~ U(0, 1), map into the truncated CDF
-            # interval [F(a), F(b)], then invert through the normal quantile.
-            cdf_a = 0.0 if lo is None else ndtr(a)
-            cdf_b = 1.0 if hi is None else ndtr(b)
-            u_samples = rng.uniform(0.0, 1.0, size)
-            p = cdf_a + u_samples * (cdf_b - cdf_a)
-            # Keep the quantile finite (avoids +-inf at the open ends).
-            p = jnp.clip(p, 1e-7, 1.0 - 1e-7)
-            samples = mean + std * ndtri(p)
+        # Vectorized inverse-CDF sampling that also handles a degenerate (std == 0)
+        # distribution element-wise. A scalar ``if std == 0`` would raise
+        # ("truth value ambiguous") for array-valued ``std`` (a per-element
+        # parametrization that the sibling distributions accept), so instead use a
+        # "safe" denominator: where ``std == 0`` the term ``std * ndtri(p)`` is 0,
+        # leaving ``samples == mean`` (then clamped into ``[low, high]`` below).
+        mean_arr = jnp.asarray(mean, dtype=float)
+        std_arr = jnp.asarray(std, dtype=float)
+        safe_std = jnp.where(std_arr == 0, 1.0, std_arr)
+
+        a = -jnp.inf if lo is None else (lo - mean_arr) / safe_std
+        b = jnp.inf if hi is None else (hi - mean_arr) / safe_std
+        # Inverse-CDF sampling: draw u ~ U(0, 1), map into the truncated CDF
+        # interval [F(a), F(b)], then invert through the normal quantile.
+        cdf_a = 0.0 if lo is None else ndtr(a)
+        cdf_b = 1.0 if hi is None else ndtr(b)
+        u_samples = rng.uniform(0.0, 1.0, size)
+        p = cdf_a + u_samples * (cdf_b - cdf_a)
+        # Keep the quantile finite (avoids +-inf at the open ends).
+        p = jnp.clip(p, 1e-7, 1.0 - 1e-7)
+        samples = mean_arr + std_arr * ndtri(p)
 
         if lo is not None or hi is not None:
             samples = jnp.clip(