Advanced Operators Tutorial

Metadata	Value
Level	Intermediate
Runtime	~30 min
Prerequisites	Operators Tutorial
Format	Python + Jupyter

Overview

Master advanced operators in Datarax: probabilistic application, random operator selection, and spatial dropout. These operators enable sophisticated augmentation pipelines and AutoAugment-style learned policies.

Learning Goals

By the end of this tutorial, you will be able to:

1. Apply transformations probabilistically with ProbabilisticOperator
2. Randomly select from multiple operators with SelectorOperator
3. Use patch-based occlusion with PatchDropoutOperator
4. Build AutoAugment-style augmentation pipelines
5. Understand the JAX compatibility patterns used

Coming from PyTorch?

PyTorch torchvision	Datarax
transforms.RandomApply([t], p=0.5)	ProbabilisticOperator(config, ...)
transforms.RandomChoice([t1, t2])	SelectorOperator(config, ...)
transforms.RandomErasing()	PatchDropoutOperator(config, ...)

Coming from TensorFlow?

TensorFlow	Datarax
Custom with tf.random.uniform < p	ProbabilisticOperator
Custom with tf.cond	SelectorOperator
tf.image.random_erasing (TF Addons)	PatchDropoutOperator

Files

• Python Script: examples/core/09_advanced_operators_tutorial.py
• Jupyter Notebook: examples/core/09_advanced_operators_tutorial.ipynb

Quick Start

Run the Python Script

python examples/core/09_advanced_operators_tutorial.py

Run the Jupyter Notebook

jupyter lab examples/core/09_advanced_operators_tutorial.ipynb

Key Concepts

ProbabilisticOperator

Wraps any operator and applies it with a configured probability.

from datarax.operators.probabilistic_operator import (
    ProbabilisticOperator,
    ProbabilisticOperatorConfig,
)

# Wrap operator with 50% probability
prob_brightness = ProbabilisticOperator(
    ProbabilisticOperatorConfig(
        operator=brightness_op,
        probability=0.5,
    ),
    rngs=nnx.Rngs(augment=42),
)

Behavior by Probability:

Probability	Mode	Behavior
p = 0.0	Deterministic	Never apply (passthrough)
0 < p < 1	Stochastic	Apply with probability p
p = 1.0	Deterministic	Always apply

Terminal Output:

Effect of probability values:
  p=0.00 (deterministic): mean delta = -0.0000
  p=0.25 (stochastic  ): mean delta = +0.0504
  p=0.50 (stochastic  ): mean delta = +0.0953
  p=0.75 (stochastic  ): mean delta = +0.1475
  p=1.00 (deterministic): mean delta = +0.1800

SelectorOperator

Randomly selects ONE operator from a list to apply per sample.

from datarax.operators.selector_operator import (
    SelectorOperator,
    SelectorOperatorConfig,
)

selector = SelectorOperator(
    SelectorOperatorConfig(
        operators=[op_bright, op_contrast, op_noise],
        weights=[0.5, 0.3, 0.2],  # 50%, 30%, 20%
    ),
    rngs=nnx.Rngs(augment=100),
)

Terminal Output:

SelectorOperator created:
  Operators: [Brightness, Contrast, Noise]
  Weights: [50%, 30%, 20%]
  Always stochastic: True

PatchDropoutOperator

Applies patch-based occlusion by dropping random rectangular regions.

from datarax.operators.modality.image.patch_dropout_operator import (
    PatchDropoutOperator,
    PatchDropoutOperatorConfig,
)

patch_dropout = PatchDropoutOperator(
    PatchDropoutOperatorConfig(
        field_key="image",
        num_patches=4,
        patch_size=(8, 8),
        drop_value=0.0,  # Black patches
        stochastic=True,
        stream_name="patch",
    ),
    rngs=nnx.Rngs(patch=42),
)

Configuration:

Parameter	Description	Default
num_patches	Number of patches to drop	4
patch_size	Size as (height, width)	(8, 8)
drop_value	Fill value for dropped regions	0.0

Building AutoAugment-Style Pipelines

Combine operators to build sophisticated augmentation pipelines:

# 1. Probabilistically apply brightness (60%)
prob_bright = ProbabilisticOperator(
    ProbabilisticOperatorConfig(operator=bright, probability=0.6, stream_name="augment"),
    rngs=nnx.Rngs(augment=100),
)

# 2. Probabilistically apply contrast (60%)
prob_contrast = ProbabilisticOperator(
    ProbabilisticOperatorConfig(operator=contrast, probability=0.6, stream_name="augment"),
    rngs=nnx.Rngs(augment=200),
)

# 3. Randomly select: noise (70%) or patch dropout (30%)
final_selector = SelectorOperator(
    SelectorOperatorConfig(
        operators=[noise, patch],
        weights=[0.7, 0.3],
    ),
    rngs=nnx.Rngs(augment=300),
)

# Build pipeline
pipeline = (
    Pipeline(source=source, stages=[prob_bright, prob_contrast, final_selector], batch_size=32, rngs=nnx.Rngs(0))
)

JAX Compatibility Patterns

All advanced operators use specific patterns for JAX compatibility:

Pattern	Used In	Purpose
jax.lax.cond	ProbabilisticOperator	Conditional execution
jax.lax.switch	SelectorOperator	Multi-way branching
jax.lax.fori_loop	PatchDropoutOperator	Loop over patches
One PRNG key per record, passed to apply	All	vmap compatibility

Why These Patterns?

1. No Python if statements: Traced JAX values can't be used in Python conditionals
2. A key per record: apply draws only from its key argument and an operator keeps no Rngs, so nothing mutates RNG state inside vmap
3. Fixed output shapes: vmap requires consistent shapes across branches

Results

Running the tutorial produces:

============================================================
Advanced Operators Tutorial
============================================================

1. ProbabilisticOperator (p=0.5):
   Output mean: 0.5987

2. SelectorOperator (3 operators):
   Output mean: 0.5623

3. PatchDropoutOperator (4 patches):
   Output mean: 0.4521 (lower due to black patches)

============================================================
Tutorial completed successfully!
============================================================

Operator Summary

Operator	Purpose	Key Config
ProbabilisticOperator	Apply with probability	probability (0-1)
SelectorOperator	Random selection	operators, weights
PatchDropoutOperator	Spatial dropout	num_patches, patch_size

Use Cases

Use Case	Operator(s)
Stochastic augmentation	ProbabilisticOperator
AutoAugment policies	SelectorOperator + ProbabilisticOperator
Occlusion robustness	PatchDropoutOperator
Test-time augmentation	All, with various settings

Next Steps

• Composition Strategies - Combine operators
• MixUp/CutMix Tutorial - Batch mixing
• Performance Guide - Optimization

API Reference

• ProbabilisticOperator
• SelectorOperator
• PatchDropoutOperator