Refine audit categorization and documentation for dual_use and safe paths in JAX and Flax

Author: pjwerneck

packages/syft-restrict/docs/blacklist.md

@@ -271,7 +271,7 @@ disallow_functions=[
     # multi-host / distributed — outbound network connections
     "jax.distributed.*", "jax.monitoring.*", "jax.experimental.multihost_utils.*",
     # FFI / interop
-    "jax.dlpack.*", "jax.ffi*",
+    "jax.dlpack.*", "jax.ffi.*",
     # array <-> disk (save/savez/savez_compressed/savetxt, load/loadtxt, tofile/fromfile/memmap)
     "jax.numpy.save*", "jax.numpy.load*", "*.tofile", "*.fromfile", "*.memmap",
     # checkpointing / serialization (write to disk, incl. gs://)

packages/syft-restrict/src/syft_restrict/audit.py

@@ -11,11 +11,11 @@
 
 - ``"unsafe"`` — matches a catalog entry for known disk/network/host-callback surface, OR is a glob
   (``jax.*``) that grants a whole namespace. Remove it or tighten the allow.
-- ``"dual_use"`` — a useful, mostly-safe op that can still be abused in combination (e.g. ``einsum``,
-  ``softmax``, ``where``). Allowed, but flagged: the *category itself* carries the "handle with care"
-  signal, so entry notes stay terse and vague — never an abuse how-to.
-- ``"safe"`` — matches a curated entry for a genuinely inert path (constants, masks, module refs) with
-  no residual output channel of its own.
+- ``"safe"`` — pure computation: ordinary math (``einsum``, ``matmul``, activations, reductions,
+  comparisons, reshapes), constants, RNG, and initializers.
+- ``"dual_use"`` — a path flagged for a specific capability beyond pure computation; each entry
+  states its own concrete reason (e.g. crossing the host/device boundary, or widening what the
+  verifier accepts as valid attribute access).
 - ``"review"`` — none of the above. The audit makes **no** guess about it: it is reported as
   uncatalogued and deferred to human review. Unknowns are never assumed safe.
 

packages/syft-restrict/src/syft_restrict/catalog/README.md

@@ -20,8 +20,7 @@ catalog/
   `0.19`. **There is no version-agnostic fallback per library:** if no version dir matches the
   installed version, that library contributes no rules and its paths fall to `review`. Add a version
   dir to cover a release.
-- `_common/default/catalog.json` is always merged in (its `default` segment is a fixed name, not a
-  version). It holds truly cross-library patterns (`*.io_callback`, `*.tofile`, …) and blanket rules
+- `_common/default/catalog.json` is always merged in. It holds truly cross-library patterns (`*.io_callback`, `*.tofile`, …) and blanket rules
   for libraries whose import root cannot be version-keyed (e.g. `orbax`: the `orbax` import root is a
   namespace package with no `__version__`; the distribution is `orbax-checkpoint`).
 
@@ -33,16 +32,21 @@ Each `catalog.json` is:
 {
   "_about": "free-text note",
   "unsafe":   { "<dotted-path glob>": "why it is unsafe" },
-  "dual_use": { "<dotted-path glob>": "what the op is (terse; the category carries the caution)" },
-  "safe":     { "<dotted-path glob>": "why it is genuinely inert" }
+  "dual_use": { "<dotted-path glob>": "the concrete reason it is flagged" },
+  "safe":     { "<dotted-path glob>": "what the op is (terse)" }
 }
 ```
 
 - `unsafe` = known disk/network/host-callback surface.
-- `dual_use` = a useful op that is mostly safe but can be **abused in combination** (`einsum`,
-  `softmax`, `where`, …). Allowed but flagged. The *category* is the caution, so keep each note a
-  terse description of what the op is — do **not** spell out abuse mechanics (no how-to).
-- `safe` = genuinely inert (constants, masks, module refs) with no residual channel of its own.
+- `safe` = pure computation: ordinary math (`einsum`, `matmul`, activations, reductions, reshapes),
+  constants, RNG, and initializers.
+- `dual_use` = a path flagged for a specific capability beyond pure computation. Each entry must
+  state its own concrete reason (e.g. crossing the host/device boundary).
 
 A path is matched strictest-first (`unsafe` → `dual_use` → `safe`). Anything matched by none defaults
 to `review` — never silently to `safe`.
+
+> [!Note]
+>
+> `safe` means "no disk/network/host-callback capability", not "no information
+> flow". The catalog lists capabilities, not guarantees.

packages/syft-restrict/src/syft_restrict/catalog/_common/default/catalog.json

@@ -2,7 +2,6 @@
   "_about": "Library-agnostic risk rules, merged into every path's audit regardless of library or version. Patterns are dotted-path globs (fnmatch); 'unsafe' = disk/network/host-callback. Also holds blanket rules for libraries whose import root cannot be version-keyed (e.g. orbax: the 'orbax' import root is a namespace package with no __version__; the distribution is orbax-checkpoint). This is advisory, not a proof.",
   "unsafe": {
     "*.fromfile": "Reads raw bytes from a file on disk.",
-    "*.host_callback*": "Legacy host-callback API — runs host Python.",
     "*.io_callback": "Runs a host Python callback with side effects (I/O).",
     "*.memmap": "Memory-maps a file on disk.",
     "*.tofile": "Writes raw array bytes to a file on disk.",

packages/syft-restrict/src/syft_restrict/catalog/flax/0.12/catalog.json

@@ -1,9 +1,104 @@
 {
-  "_about": "Risk rules for flax 0.12.x. 'unsafe' = disk/network/host-callback; 'dual_use' = useful surface that is mostly safe but can be abused in combination (the category itself is the caution, so notes stay terse and vague); 'safe' = genuinely inert. Cross-library rules in _common/default are merged in on top of these. This is advisory, not a proof.",
+  "_about": "Risk rules for flax 0.12.x. 'unsafe' = disk/network/host-callback; 'safe' = pure computation; 'dual_use' = flagged for a specific capability beyond pure computation, reason stated per entry. Covers the commonly-used flax.linen surface; anything uncatalogued falls through to 'review'. Cross-library rules in _common/default are merged in on top of these. This is advisory, not a proof.",
   "dual_use": {
-    "flax.linen.Module": "Flax module base class; required to define a model. Keep allow_base_class_attributes=False to require explicit assignment."
+    "flax.linen.Module": "Flax module base class; required to define a model. allow_base_class_attributes widens what the verifier treats as valid attribute access — keep it False to require explicit assignment."
+  },
+  "safe": {
+    "flax.linen.BatchNorm": "Batch normalization.",
+    "flax.linen.Bidirectional": "Bidirectional wrapper running a cell forward and backward.",
+    "flax.linen.Conv": "Convolution layer.",
+    "flax.linen.ConvLSTMCell": "Convolutional LSTM cell.",
+    "flax.linen.ConvLocal": "Locally-connected (unshared-weight) convolution layer.",
+    "flax.linen.ConvTranspose": "Transposed (fractionally-strided) convolution layer.",
+    "flax.linen.Dense": "Fully-connected (dense) layer.",
+    "flax.linen.DenseGeneral": "Dense layer contracting over arbitrary axes.",
+    "flax.linen.Dropout": "Stochastic dropout layer.",
+    "flax.linen.Einsum": "Einsum layer with a learned operand.",
+    "flax.linen.Embed": "Embedding lookup layer.",
+    "flax.linen.GRUCell": "GRU recurrent cell.",
+    "flax.linen.GroupNorm": "Group normalization.",
+    "flax.linen.InstanceNorm": "Instance normalization.",
+    "flax.linen.LSTMCell": "LSTM recurrent cell.",
+    "flax.linen.LayerNorm": "Layer normalization.",
+    "flax.linen.MultiHeadAttention": "Multi-head attention layer.",
+    "flax.linen.MultiHeadDotProductAttention": "Multi-head dot-product attention layer.",
+    "flax.linen.OptimizedLSTMCell": "LSTM cell with a fused implementation.",
+    "flax.linen.PReLU": "Parametric-ReLU activation layer.",
+    "flax.linen.Partitioned": "Sharding-metadata wrapper for a parameter; structural annotation.",
+    "flax.linen.RMSNorm": "Root-mean-square layer normalization.",
+    "flax.linen.RNN": "Recurrent layer applying a cell over a sequence.",
+    "flax.linen.SelfAttention": "Multi-head self-attention layer.",
+    "flax.linen.Sequential": "Container applying sub-layers in order; the sub-layers are allow-listed separately.",
+    "flax.linen.SimpleCell": "Elman (simple) recurrent cell.",
+    "flax.linen.SpectralNorm": "Spectral-normalization wrapper.",
+    "flax.linen.Variable": "Handle to a module variable; a state reference, not compute or IO.",
+    "flax.linen.WeightNorm": "Weight-normalization wrapper.",
+    "flax.linen.avg_pool": "Average pooling over windows.",
+    "flax.linen.celu": "CELU activation.",
+    "flax.linen.checkpoint": "Rematerialization transform (alias of remat); not disk checkpointing.",
+    "flax.linen.combine_masks": "Combines boolean masks with logical AND.",
+    "flax.linen.compact": "Decorator marking a module's inline-submodule method; structural.",
+    "flax.linen.cond": "Lifted conditional executing one of two branches.",
+    "flax.linen.dot_product_attention": "Dot-product attention function.",
+    "flax.linen.dot_product_attention_weights": "Dot-product attention weights (softmax over scores).",
+    "flax.linen.elu": "ELU activation.",
+    "flax.linen.gelu": "GELU activation.",
+    "flax.linen.glu": "Gated-linear-unit activation.",
+    "flax.linen.hard_sigmoid": "Hard-sigmoid activation.",
+    "flax.linen.hard_silu": "Hard-SiLU / hard-swish activation.",
+    "flax.linen.hard_swish": "Hard-swish / hard-SiLU activation.",
+    "flax.linen.hard_tanh": "Hard-tanh activation.",
+    "flax.linen.initializers.constant": "Constant-value parameter initializer.",
+    "flax.linen.initializers.delta_orthogonal": "Delta-orthogonal initializer for convolution kernels.",
+    "flax.linen.initializers.glorot_normal": "Glorot (Xavier) normal initializer.",
+    "flax.linen.initializers.glorot_uniform": "Glorot (Xavier) uniform initializer.",
+    "flax.linen.initializers.he_normal": "He (Kaiming) normal initializer.",
+    "flax.linen.initializers.he_uniform": "He (Kaiming) uniform initializer.",
+    "flax.linen.initializers.kaiming_normal": "Kaiming (He) normal initializer.",
+    "flax.linen.initializers.kaiming_uniform": "Kaiming (He) uniform initializer.",
+    "flax.linen.initializers.lecun_normal": "LeCun normal initializer.",
+    "flax.linen.initializers.lecun_uniform": "LeCun uniform initializer.",
+    "flax.linen.initializers.normal": "Normal-distribution parameter initializer.",
+    "flax.linen.initializers.ones": "All-ones parameter initializer.",
+    "flax.linen.initializers.orthogonal": "Orthogonal-matrix parameter initializer.",
+    "flax.linen.initializers.truncated_normal": "Truncated-normal parameter initializer.",
+    "flax.linen.initializers.uniform": "Uniform-distribution parameter initializer.",
+    "flax.linen.initializers.variance_scaling": "Variance-scaling parameter initializer.",
+    "flax.linen.initializers.xavier_normal": "Xavier (Glorot) normal initializer.",
+    "flax.linen.initializers.xavier_uniform": "Xavier (Glorot) uniform initializer.",
+    "flax.linen.initializers.zeros": "All-zeros parameter initializer.",
+    "flax.linen.jit": "Lifted JIT compilation of a module.",
+    "flax.linen.leaky_relu": "Leaky-ReLU activation.",
+    "flax.linen.log_sigmoid": "Log-sigmoid activation.",
+    "flax.linen.log_softmax": "Log-softmax activation.",
+    "flax.linen.logsumexp": "Log-sum-exp reduction.",
+    "flax.linen.make_attention_mask": "Builds an attention mask by broadcasting a query/key position comparison.",
+    "flax.linen.make_causal_mask": "Builds a causal (lower-triangular) attention mask from a shape.",
+    "flax.linen.map_variables": "Transform mapping over a module's variable collections.",
+    "flax.linen.max_pool": "Max pooling over windows.",
+    "flax.linen.nowrap": "Decorator opting a method out of module auto-wrapping; structural.",
+    "flax.linen.one_hot": "One-hot encoding of integer indices.",
+    "flax.linen.pool": "Generic windowed pooling reduction.",
+    "flax.linen.relu": "ReLU activation.",
+    "flax.linen.relu6": "ReLU6 activation.",
+    "flax.linen.remat": "Rematerialization (gradient-checkpointing) transform.",
+    "flax.linen.remat_scan": "Rematerialized scan transform.",
+    "flax.linen.scan": "Lifted scan over a module along an axis.",
+    "flax.linen.selu": "SELU activation.",
+    "flax.linen.sigmoid": "Sigmoid activation.",
+    "flax.linen.silu": "SiLU / swish activation.",
+    "flax.linen.soft_sign": "Soft-sign activation.",
+    "flax.linen.softmax": "Softmax activation.",
+    "flax.linen.softplus": "Softplus activation.",
+    "flax.linen.standardize": "Standardizes to zero-mean, unit-variance along an axis.",
+    "flax.linen.swish": "Swish / SiLU activation.",
+    "flax.linen.switch": "Lifted multi-branch switch.",
+    "flax.linen.tanh": "Hyperbolic-tangent activation.",
+    "flax.linen.vmap": "Lifted vectorizing map over a module.",
+    "flax.linen.while_loop": "Lifted while loop."
   },
   "unsafe": {
+    "flax.io.*": "File IO shim — reads/writes files on disk (and gs:// cloud storage when available).",
     "flax.serialization.*": "Serializes model state to bytes / disk.",
     "flax.training.checkpoints.*": "Writes / reads checkpoints to disk (path accepts gs://)."
   }