dfax 0.1.0__py3-none-any.whl

dfax/__init__.py

@@ -0,0 +1,2 @@
+from dfax.dfax import *
+from dfax.utils import *

dfax/dfax.py

@@ -0,0 +1,356 @@
+import jax
+import chex
+import jax.numpy as jnp
+from flax import struct
+
+
+@jax.jit
+def DFAx(start, transitions, labels):
+    n_states, n_tokens = transitions.shape
+    trans_flat = transitions.flatten()
+    is_reach_init = jnp.zeros((n_states,), dtype=bool).at[start].set(True)
+
+    def step(is_reach: jnp.ndarray) -> jnp.ndarray:
+        reach_repeat = jnp.repeat(is_reach, n_tokens)
+        dest_counts = jnp.zeros((n_states,), dtype=jnp.int32).at[trans_flat].add(reach_repeat)
+        return dest_counts > 0
+
+    def cond(pair):
+        prev_is_reach, curr_is_reach = pair
+        return jnp.any(prev_is_reach != curr_is_reach)
+
+    def body(pair):
+        prev_is_reach, curr_is_reach = pair
+        next_is_reach = step(curr_is_reach)
+        return (curr_is_reach, next_is_reach)
+
+    is_reach, _ = jax.lax.while_loop(cond, body, (is_reach_init, step(is_reach_init)))
+
+    return _DFAx(start=start,
+                 transitions=transitions,
+                 labels=labels,
+                 is_reach=is_reach)
+
+
+@struct.dataclass
+class _DFAx:
+    start: int
+    transitions: jnp.ndarray
+    labels: jnp.ndarray
+    is_reach: jnp.ndarray
+
+    @property
+    def n_states(self):
+        return jnp.sum(self.is_reach)
+
+    @property
+    def max_n_states(self):
+        return self.transitions.shape[0]
+
+    @property
+    def n_tokens(self):
+        return self.transitions.shape[1]
+
+    @property
+    def is_reach_tile(self) -> jnp.ndarray:
+        return jnp.tile(self.is_reach.reshape(-1, 1), (1, self.n_tokens))
+
+    @jax.jit
+    def __eq__(self, other: "DFAx") -> jnp.ndarray:
+        n = min(self.max_n_states, other.max_n_states)
+
+        start_eq = (self.start == other.start)
+        transitions_eq = jnp.where(
+            self.n_tokens == other.n_tokens,
+            jnp.all(self.transitions[:n] == other.transitions[:n]),
+            False
+        )
+        labels_eq = jnp.all(self.labels[:n] == other.labels[:n])
+
+        return jnp.logical_and(start_eq, jnp.logical_and(transitions_eq, labels_eq))
+
+
+    @jax.jit
+    def advance(self, symbol: int) -> "DFAx":
+        return DFAx(
+            start=jnp.where(
+                jnp.logical_and(symbol >= 0, symbol < self.n_tokens),
+                self.transitions[self.start, symbol],
+                self.start
+            ),
+            transitions=self.transitions,
+            labels=self.labels)
+
+    @jax.jit
+    def mutate(self, key: chex.PRNGKey) -> "DFAx":
+        key, k1, k2 = jax.random.split(key, 3)
+
+        flat_is_reach_tile = self.is_reach_tile.flatten()
+
+        scores, indices = jax.lax.top_k(jnp.where(flat_is_reach_tile, 1, 0), flat_is_reach_tile.size)
+        num_valid = jnp.sum(scores > 0)
+        valid_idx = jax.random.randint(k1, (), 0, num_valid)
+        flat_idx = indices[valid_idx]
+        s, a = jnp.divmod(flat_idx, self.n_tokens)
+
+        scores, indices = jax.lax.top_k(jnp.where(self.is_reach, 1, 0), self.max_n_states)
+        num_valid = jnp.sum(scores > 0)
+        valid_idx = jax.random.randint(k2, (), 0, num_valid)
+        t = indices[valid_idx]
+
+        transitions = self.transitions.at[s, a].set(t)
+
+        return DFAx(start=self.start,
+                    transitions=transitions,
+                    labels=self.labels)
+
+    @jax.jit
+    def mutate_reject_lang(self, key: chex.PRNGKey) -> "DFAx":
+        key, k1, k2 = jax.random.split(key, 3)
+
+        is_self_loop = self.transitions == jnp.arange(self.max_n_states)[:, None]
+        sa_mask = jnp.logical_and(
+            jnp.logical_and(self.is_reach_tile, is_self_loop),
+            jnp.logical_not(self.labels)[:, None]
+        ).flatten()
+
+        scores, indices = jax.lax.top_k(jnp.where(sa_mask, 1, 0), sa_mask.size)
+        num_valid = jnp.sum(scores > 0)
+        valid_idx = jax.random.randint(k1, (), 0, num_valid)
+        flat_idx = indices[valid_idx]
+        s, a = jnp.divmod(flat_idx, self.n_tokens)
+
+        is_sink = jnp.all(is_self_loop, axis=-1)
+        is_reject = jnp.logical_and.reduce(jnp.array([self.is_reach, is_sink, jnp.logical_not(self.labels)]))
+        t_mask = jnp.where(jnp.any(is_reject), is_reject, jnp.logical_not(self.is_reach))
+
+        scores, indices = jax.lax.top_k(jnp.where(t_mask, 1, 0), t_mask.size)
+        num_valid = jnp.sum(scores > 0)
+        valid_idx = jax.random.randint(k2, (), 0, num_valid)
+        t = indices[valid_idx]
+
+        transitions = self.transitions.at[s, a].set(t)
+
+        return DFAx(start=self.start,
+                    transitions=transitions,
+                    labels=self.labels)
+
+    @jax.jit
+    def sink_accepts(self) -> "DFAx":
+        accept_mask = jnp.tile(self.labels.reshape(-1, 1), (1, self.n_tokens))
+        replace_mask = accept_mask & self.is_reach_tile
+
+        sink_indices = jnp.tile(jnp.arange(self.max_n_states).reshape(-1, 1), (1, self.n_tokens))
+        transitions = jnp.where(replace_mask, sink_indices, self.transitions)
+
+        return DFAx(
+            start=self.start,
+            transitions=transitions,
+            labels=self.labels
+        )
+
+    @jax.jit
+    def prune(self) -> "DFAx":
+        pruned_transitions = jnp.where(self.is_reach_tile, self.transitions, jnp.arange(self.max_n_states)[:, None])
+        pruned_labels = jnp.where(self.is_reach, self.labels, False)
+
+        return DFAx(start=self.start,
+                    transitions=pruned_transitions,
+                    labels=pruned_labels)
+
+    @jax.jit
+    def minimize(self) -> "DFAx":
+        return self.naivePR().prune().canonicalize()
+
+    @jax.jit
+    def naivePR(self) -> "DFAx":
+        # Algorithm 2 from https://arxiv.org/pdf/2410.22764
+        q_f = jnp.argmax(jnp.where(self.is_reach, self.labels, 0))
+        q_n = jnp.argmin(jnp.where(self.is_reach, self.labels, 1))
+        block = jnp.where(self.labels, q_f, q_n)
+        block = jnp.where(self.is_reach, block, jnp.arange(self.max_n_states))
+
+        qs = jnp.arange(self.max_n_states)
+        as_ = jnp.arange(self.n_tokens)
+        qas = jnp.stack(jnp.meshgrid(qs, as_, indexing="ij"), axis=-1).reshape(-1, 2)
+
+        def iteration(state):
+            block, _ = state
+
+            def elect(q_a):
+                q, a = q_a
+                return jax.lax.cond(
+                    self.is_reach[q],
+                    lambda _: (
+                        block[q],
+                        jnp.where(
+                            block[self.transitions[q, a]] != block[self.transitions[block[q], a]],
+                            q,
+                            -1
+                        )
+                    ),
+                    lambda _: (block[q], -1),
+                    operand=None
+                )
+            blk_idx, leaders = jax.vmap(elect)(qas)
+            new_leader = jax.ops.segment_max(leaders, blk_idx, num_segments=self.max_n_states)
+
+            def assign(q_a):
+                q, a = q_a
+                return jax.lax.cond(
+                    self.is_reach[q],
+                    lambda _: (
+                        q,
+                        jnp.where(
+                            (block[self.transitions[q, a]] != block[self.transitions[block[q], a]]),
+                            new_leader[block[q]],
+                            -1
+                        )
+                    ),
+                    lambda _: (q, -1),
+                    operand=None
+                )
+            qs_idx, new_vals = jax.vmap(assign)(qas)
+            new_block = jax.ops.segment_max(new_vals, qs_idx, num_segments=self.max_n_states)
+            new_block = jnp.where(new_block < 0, block, new_block)
+
+            return (new_block, jnp.any(new_block != block))
+
+        block, _ = jax.lax.while_loop(
+            lambda s: s[1],
+            lambda s: iteration((s[0], False)),
+            (block, True)
+        )
+
+        minimized_start       = block[self.start]
+        minimized_labels      = self.labels[block]
+        minimized_transitions = block[self.transitions]
+
+        return DFAx(start=minimized_start,
+                    transitions=minimized_transitions,
+                    labels=minimized_labels)
+
+    @jax.jit
+    def canonicalize(self) -> "DFAx":
+        old_to_new = (-jnp.ones((self.max_n_states,), dtype=jnp.int32)).at[self.start].set(0)
+        visited    = jnp.zeros((self.max_n_states,), dtype=bool).at[self.start].set(True)
+        queue      = (-jnp.ones((self.max_n_states,), dtype=jnp.int32)).at[0].set(self.start)
+        head       = 0
+        tail       = (head + 1) % self.max_n_states
+        count      = 0
+
+        def cond(carry):
+            _, _, _, head, tail, _ = carry
+            return head != tail
+
+        def body(carry):
+            visited, old_to_new, queue, head, tail, count = carry
+
+            current_state = queue[head]
+            head = (head + 1) % self.max_n_states
+
+            old_to_new = old_to_new.at[current_state].set(count)
+            count += 1
+
+            next_states = self.transitions[current_state]
+
+            def push(i, carry):
+                visited, queue, tail = carry
+                ns = next_states[i]
+                unseen = jnp.logical_not(visited[ns])
+
+                queue = queue.at[tail].set(jnp.where(unseen, ns, queue[tail]))
+                tail  = (tail + unseen) % self.max_n_states
+
+                visited = visited.at[ns].set(True)
+                return visited, queue, tail
+
+            visited, queue, tail = jax.lax.fori_loop(
+                0, self.n_tokens, push, (visited, queue, tail)
+            )
+
+            return visited, old_to_new, queue, head, tail, count
+
+        visited, old_to_new, queue, head, tail, count = jax.lax.while_loop(
+            cond, body, (visited, old_to_new, queue, head, tail, count)
+        )
+
+        mask = (old_to_new < 0)
+        ranks = jnp.cumsum(mask) - 1
+        fill_vals = count + ranks
+        old_to_new = jnp.where(mask, fill_vals, old_to_new)
+
+        start        = old_to_new[self.start]
+        transitions  = self.transitions.at[old_to_new].set(old_to_new[self.transitions])
+        labels       = self.labels.at[old_to_new].set(self.labels)
+
+        return DFAx(start=start,
+                    transitions=transitions,
+                    labels=labels)
+
+    @jax.jit
+    def to_graph(self):
+        srcs, tgts = jnp.meshgrid(jnp.arange(self.max_n_states), jnp.arange(self.max_n_states), indexing="ij")
+        srcs = srcs.flatten()
+        tgts = tgts.flatten()
+
+        edge_index = jnp.stack([srcs, tgts])
+        
+        is_init = jnp.logical_and(
+            self.is_reach,
+            jnp.arange(self.max_n_states) == self.start
+        )
+        is_accept = jnp.logical_and(
+            self.is_reach,
+            self.labels
+        )
+        is_reject = jnp.logical_and(
+            self.is_reach,
+            jnp.all(
+                jnp.logical_and(
+                    self.transitions == jnp.arange(self.max_n_states)[:, None], # has all loops
+                    jnp.logical_not(self.labels[:, None]) # not accepting
+                ),
+                axis=1
+            )
+        )
+        is_non_terminal = jnp.logical_and(
+            self.is_reach,
+            jnp.logical_and(
+                jnp.logical_not(is_accept),
+                jnp.logical_not(is_reject)
+            )
+        )
+
+        node_features = jnp.stack([is_init, is_accept, is_reject, is_non_terminal], axis=1).astype(jnp.float32)
+
+        edge_features = jnp.logical_and(
+            self.is_reach[:, None, None],
+            self.transitions[:, None, :] == jnp.arange(self.max_n_states)[None, :, None] # one hot tokens
+        ).astype(jnp.float32).reshape(-1, self.n_tokens)
+
+        mask = jnp.any(edge_features != 0, axis=-1)[:, None]
+        edge_features = jnp.concatenate([node_features[srcs] * mask, edge_features, node_features[tgts] * mask], axis=-1)
+
+        graph = {
+            "node_features": node_features,
+            "edge_features": edge_features,
+            "edge_index": edge_index,
+            "current_state": jnp.array([self.start]),
+            "n_states": jnp.full(self.max_n_states, self.n_states)
+        }
+
+        return graph
+
+    @jax.jit
+    def reward(self, binary: bool = False) -> float:
+        is_accept = self.labels[self.start]
+        start_row = self.transitions[self.start]
+        start_vec = jnp.full((self.n_tokens,), self.start, dtype=start_row.dtype)
+        is_sink = jnp.all(start_row == start_vec)
+
+        return jnp.where(binary,
+            jnp.where(is_accept, 1.0, 0.0),
+            jnp.where(is_accept, 1.0, jnp.where(is_sink, -1.0, 0.0))
+        )
+

dfax/samplers.py

@@ -0,0 +1,179 @@
+import jax
+import chex
+import jax.numpy as jnp
+from flax import struct
+from dfax import DFAx
+from functools import partial
+
+
+# Base sampler: holds parameters
+@struct.dataclass
+class DFASampler:
+    n_tokens: int = 10
+    max_size: int = 10
+    p: float | None = None
+
+    @partial(jax.jit, static_argnums=(0,))
+    def sample(self, key: chex.PRNGKey) -> DFAx:
+        raise NotImplementedError
+
+    @partial(jax.jit, static_argnums=(0, 2))
+    def sample_n(self, key: chex.PRNGKey, lower_bound: int = 2):
+        if self.p is not None:
+            values = jnp.arange(lower_bound, self.max_size + 1)
+            weights = self.p ** values
+            weights = weights / jnp.sum(weights)
+            idx = jax.random.choice(key, values, p=weights)
+            return idx
+        else:
+            return jax.random.randint(key, (), lower_bound, self.max_size + 1)
+
+    @partial(jax.jit, static_argnums=(0,))
+    def trivial(self, label):
+        start = 0
+        transitions = jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens))
+        labels = jnp.zeros((self.max_size,), dtype=bool).at[start].set(label)
+        return DFAx(start=start,
+                     transitions=transitions,
+                     labels=labels)
+
+
+# Reach sampler
+@struct.dataclass
+class ReachSampler(DFASampler):
+    prob_stutter: float = 0.9
+
+    @partial(jax.jit, static_argnums=(0,))
+    def sample(self, key: chex.PRNGKey) -> DFAx:
+        key, subkey = jax.random.split(key)
+        n = self.sample_n(subkey, lower_bound=2)
+        success = n-1
+        transitions = jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens))
+        labels = jnp.zeros(self.max_size, dtype=bool)
+        labels = labels.at[success].set(True)
+        transitions = transitions.at[success, :].set(success)
+
+        def body_fn(i, carry):
+            transitions, labels, key = carry
+            key, k1, k2, k3 = jax.random.split(key, 4)
+            perm = jax.random.permutation(k1, jnp.arange(self.n_tokens))
+            row = jnp.full(self.n_tokens, i, dtype=jnp.int32)
+            row = row.at[perm[0]].set(i+1)
+            rest = perm[1:]
+            r = jax.random.uniform(k2, (self.n_tokens-1,))
+            choice = jax.random.bernoulli(k3, 0.5, (self.n_tokens-1,))
+            dest = jnp.where(r <= self.prob_stutter, i, i+1)
+            row = row.at[rest].set(dest)
+            transitions = transitions.at[i].set(row)
+            return (transitions, labels, key)
+
+        transitions, labels, _ = jax.lax.cond(
+            n == 0,
+            lambda _: (jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens)), jnp.zeros(self.max_size, dtype=bool).at[0].set(True), key),
+            lambda _: jax.lax.fori_loop(0, n-1, body_fn, (transitions, labels, key)),
+            operand=None
+        )
+        return DFAx(start=0, transitions=transitions, labels=labels).minimize()
+
+
+# Reach-Avoid sampler
+@struct.dataclass
+class ReachAvoidSampler(DFASampler):
+    prob_stutter: float = 0.9
+
+    @partial(jax.jit, static_argnums=(0,))
+    def sample(self, key: chex.PRNGKey) -> DFAx:
+        key, subkey = jax.random.split(key)
+        n = self.sample_n(subkey, lower_bound=3)
+        success, fail = n-2, n-1
+        transitions = jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens))
+        labels = jnp.zeros(self.max_size, dtype=bool)
+        labels = labels.at[success].set(True)
+        transitions = transitions.at[success, :].set(success)
+        transitions = transitions.at[fail, :].set(fail)
+
+        def body_fn(i, carry):
+            transitions, labels, key = carry
+            key, k1, k2, k3 = jax.random.split(key, 4)
+            perm = jax.random.permutation(k1, jnp.arange(self.n_tokens))
+            row = jnp.full(self.n_tokens, i, dtype=jnp.int32)
+            row = row.at[perm[0]].set(i+1)
+            row = row.at[perm[1]].set(fail)
+            rest = perm[2:]
+            r = jax.random.uniform(k2, (self.n_tokens-2,))
+            choice = jax.random.bernoulli(k3, 0.5, (self.n_tokens-2,))
+            dest = jnp.where(r <= self.prob_stutter, i, jnp.where(choice, i+1, fail))
+            row = row.at[rest].set(dest)
+            transitions = transitions.at[i].set(row)
+            return (transitions, labels, key)
+
+        transitions, labels, _ = jax.lax.cond(
+            n == 0,
+            lambda _: (jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens)), jnp.zeros(self.max_size, dtype=bool).at[0].set(True), key),
+            lambda _: jax.lax.fori_loop(0, n-2, body_fn, (transitions, labels, key)),
+            operand=None
+        )
+        return DFAx(start=0, transitions=transitions, labels=labels).minimize()
+
+
+# Reach-Avoid with random mutations
+@struct.dataclass
+class RADSampler(DFASampler):
+    p: float | None = 0.5
+    prob_stutter: float = 0.9
+    max_mutations: int = 5
+
+    @partial(jax.jit, static_argnums=(0,))
+    def sample(self, key: chex.PRNGKey) -> DFAx:
+        key, subkey = jax.random.split(key)
+        n = self.sample_n(subkey, lower_bound=3)
+        success, fail = n-2, n-1
+        transitions = jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens))
+        labels = jnp.zeros(self.max_size, dtype=bool)
+        labels = labels.at[success].set(True)
+        transitions = transitions.at[success, :].set(success)
+        transitions = transitions.at[fail, :].set(fail)
+
+        def body_fn(i, carry):
+            transitions, labels, key = carry
+            key, k1, k2, k3, k4 = jax.random.split(key, 5)
+            perm = jax.random.permutation(k1, jnp.arange(self.n_tokens))
+            row = jnp.full(self.n_tokens, i, dtype=jnp.int32)
+            row = row.at[perm[0]].set(i+1)
+            is_avoid_problem = jax.random.bernoulli(k2, 0.5)
+            row = row.at[perm[1]].set(jnp.where(is_avoid_problem, fail, i))
+            rest = perm[2:]
+            r = jax.random.uniform(k3, (self.n_tokens-2,))
+            choice = jax.random.bernoulli(k4, 0.5, (self.n_tokens-2,))
+            dest = jnp.where(r <= self.prob_stutter, i, jnp.where(choice, i+1, fail))
+            row = row.at[rest].set(dest)
+            transitions = transitions.at[i].set(row)
+            return (transitions, labels, key)
+
+        transitions, labels, _ = jax.lax.cond(
+            n == 0,
+            lambda _: (jnp.tile(jnp.arange(self.max_size).reshape(-1, 1), (1, self.n_tokens)), jnp.zeros(self.max_size, dtype=bool).at[0].set(True), key),
+            lambda _: jax.lax.fori_loop(0, n-2, body_fn, (transitions, labels, key)),
+            operand=None
+        )
+        candidate = DFAx(start=0, transitions=transitions, labels=labels).minimize()
+
+        key, subkey = jax.random.split(key)
+        n_mutations = jax.random.choice(subkey, self.max_mutations + 1)
+
+        def derive(i, carry):
+            k, cand = carry
+            k, sk = jax.random.split(k)
+            new_cand = cand.mutate(sk).sink_accepts().minimize()
+            cand = jax.lax.cond(
+                new_cand.n_states <= 1,
+                lambda _: cand,
+                lambda _: new_cand,
+                operand=None
+            )
+            return k, cand
+
+        _, candidate = jax.lax.fori_loop(0, n_mutations, derive, (key, candidate))
+
+        return candidate
+

dfax/utils.py

@@ -0,0 +1,189 @@
+import jax
+from dfa import DFA
+from dfax import DFAx
+import networkx as nx
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+
+
+def dfa2dfax(dfa: DFA) -> DFAx:
+    states = dfa.states()
+    inputs = dfa.inputs
+    start = dfa.start
+    transitions = jnp.array([[dfa._transition(s, a) for a in inputs] for s in states])
+    labels = jnp.array([dfa._label(s) for s in states])
+    tmp = DFAx(
+        start=start,
+        transitions=transitions,
+        labels=labels
+    )
+    return tmp
+
+
+def dfax2dfa(dfax: DFAx) -> DFA:
+    inputs = set(range(dfax.transitions.shape[1]))
+
+    def transition(s, a):
+        return int(dfax.transitions[s, a])
+
+    def label(s):
+        return bool(dfax.labels[s])
+
+    return DFA(
+        start=int(dfax.start),
+        inputs=inputs,
+        transition=transition,
+        label=label,
+    )
+
+
+@jax.jit
+def batch2graph(batch):
+    if batch["node_features"].ndim == 2 and batch["edge_features"].ndim == 2 and batch["edge_index"].ndim == 2 and batch["current_state"].ndim == 1:
+        return batch
+
+    batch_size, n_nodes, _ = batch["node_features"].shape
+    node_features = jnp.concatenate(batch["node_features"])
+    edge_features = jnp.concatenate(batch["edge_features"])
+    offset = (jnp.arange(batch_size, dtype=jnp.int32) * n_nodes)
+    edge_index = jnp.concatenate(batch["edge_index"] + offset[:, None, None], axis=1)
+    current_state = (batch["current_state"].reshape(batch_size, -1) + offset[:, None]).flatten()
+    n_states = jnp.concatenate(batch["n_states"])
+
+    graph = {
+        "node_features": node_features,
+        "edge_features": edge_features,
+        "edge_index": edge_index,
+        "current_state": current_state,
+        "n_states": n_states
+    }
+
+    return graph
+
+@jax.jit
+def list2batch(graphs):
+        node_features_batch = jnp.stack([graph["node_features"] for graph in graphs], axis=0)
+        node_features_batch = node_features_batch if node_features_batch.ndim == 3 else jnp.concatenate(node_features_batch, axis=0)
+
+        edge_features_batch = jnp.stack([graph["edge_features"] for graph in graphs], axis=0)
+        edge_features_batch = edge_features_batch if edge_features_batch.ndim == 3 else jnp.concatenate(edge_features_batch, axis=0)
+
+        edge_index_batch = jnp.stack([graph["edge_index"] for graph in graphs], axis=0)
+        edge_index_batch = edge_index_batch if edge_index_batch.ndim == 3 else jnp.concatenate(edge_index_batch, axis=0)
+
+        current_state_batch = jnp.stack([graph["current_state"] for graph in graphs], axis=0)
+        current_state_batch = current_state_batch if current_state_batch.ndim == 2 else jnp.concatenate(current_state_batch, axis=0)
+
+        n_states_batch = jnp.stack(jnp.array([graph["n_states"] for graph in graphs]), axis=0)
+        n_states_batch = n_states_batch if n_states_batch.ndim == 2 else jnp.concatenate(n_states_batch, axis=0)
+
+        batch = {
+            "node_features": node_features_batch,
+            "edge_features": edge_features_batch,
+            "edge_index": edge_index_batch,
+            "current_state": current_state_batch,
+            "n_states": n_states_batch,
+        }
+
+        return batch
+
+
+def visualize(dfax, label_states=False, save_path=None):
+    n_states, n_tokens = dfax.transitions.shape
+
+    G = nx.DiGraph()
+    for s in range(n_states):
+        if dfax.is_reach[s]:
+            G.add_node(s, label=str(s))
+
+    edges = {}
+    for s in range(n_states):
+        s = int(s)
+        for a in range(n_tokens):
+            a = int(a)
+            t = int(dfax.transitions[s, a])
+            if s != t:
+                if (s, t) not in edges:
+                    edges[(s, t)] = [str(a)]
+                else:
+                    edges[(s, t)].append(str(a))
+
+    for (s, t) in edges:
+        G.add_edge(s, t, label=edges[(s, t)])
+
+    # dummy start node
+    dummy_start = ""
+    G.add_node(dummy_start)
+    G.add_edge(dummy_start, int(dfax.start))
+
+    pos = nx.shell_layout(G)
+    # pos = nx.planar_layout(G)
+    start_pos = pos[int(dfax.start)]
+    pos[dummy_start] = (start_pos[0] - 0.5, start_pos[1])
+
+    accept_nodes = [s for s in G.nodes() if s != dummy_start and dfax.labels[s]]
+    reject_nodes = [s for s in G.nodes() if s != dummy_start and not dfax.labels[s] and jnp.all(dfax.transitions[s] == s)]
+    undecd_nodes = [s for s in G.nodes() if s != dummy_start and not dfax.labels[s]]
+
+    # draw nodes
+    nx.draw_networkx_nodes(G, pos, nodelist=undecd_nodes, node_size=1200,
+                           node_color="white", edgecolors="black", linewidths=2)
+    nx.draw_networkx_nodes(G, pos, nodelist=accept_nodes, node_size=1200,
+                           node_color="#88E788", edgecolors="black", linewidths=2)
+    nx.draw_networkx_nodes(G, pos, nodelist=reject_nodes, node_size=1200,
+                           node_color="#FF746C", edgecolors="black", linewidths=2)
+
+    if label_states:
+        nx.draw_networkx_labels(G, pos, font_size=20, font_weight="bold")
+
+    ax = plt.gca()
+
+    # draw edges
+    for (u, v) in G.edges():
+        if u == dummy_start:
+            nx.draw_networkx_edges(G, pos, edgelist=[(u, v)], arrows=True, arrowsize=20, node_size=1200)
+            continue
+
+        if G.has_edge(v, u):
+            rad = 0.25
+        else:
+            rad = 0.0
+
+        nx.draw_networkx_edges(G, pos, edgelist=[(u, v)], arrows=True, arrowsize=20,
+                               connectionstyle=f"arc3,rad={rad}", node_size=1200)
+
+        # --- draw tokens along the curved edge ---
+        x0, y0 = pos[u]
+        x1, y1 = pos[v]
+        n = len(edges[(u, v)])
+        for i, a in enumerate(edges[(u, v)]):
+            ratio = (i + 1) / (n + 1)
+
+            # Compute midpoint along arc3 curve
+            if rad != 0:
+                # Control point for the quadratic Bezier
+                xm_ctrl = (x0 + x1) / 2 + rad * (y1 - y0)
+                ym_ctrl = (y0 + y1) / 2 - rad * (x1 - x0)
+
+                # Quadratic Bezier formula
+                xm = (1 - ratio) ** 2 * x0 + 2 * (1 - ratio) * ratio * xm_ctrl + ratio ** 2 * x1
+                ym = (1 - ratio) ** 2 * y0 + 2 * (1 - ratio) * ratio * ym_ctrl + ratio ** 2 * y1
+            else:
+                xm = x0 * (1 - ratio) + x1 * ratio
+                ym = y0 * (1 - ratio) + y1 * ratio
+
+            circle = patches.Circle((xm, ym), 0.08, facecolor="gold", edgecolor="orange", lw=1.5, zorder=5)
+            ax.add_patch(circle)
+            ax.text(xm, ym, a, ha="center", va="center", fontsize=16, color="black", weight="bold", zorder=6)
+
+    plt.axis("equal")
+    plt.tight_layout()
+    plt.axis("off")
+
+    if save_path:
+        plt.savefig(save_path, bbox_inches="tight", dpi=300)
+    else:
+        plt.show()
+    plt.close()
+

dfax-0.1.0.dist-info/METADATA

@@ -0,0 +1,136 @@
+Metadata-Version: 2.4
+Name: dfax
+Version: 0.1.0
+Summary: Python library for modeling DFAs, Moore Machines, and Transition Systems in JAX.
+Author-email: Beyazit Yalcinkaya <beyazit@berkeley.edu>
+License-File: LICENSE
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+
+# DFAx
+
+A JAX-compatible Python implementation of a Deterministic Finite Automaton (DFA).
+
+## Installation
+
+This package will soon be made pip-installable. In the meantime, pull the repo and and install locally.
+
+```
+git clone https://github.com/rad-dfa/dfax.git
+pip install -e dfax
+```
+
+## Usage
+
+Create DFAs by specifying a `start` state, `transitions` matrix, which is max number of states by number of alphabet symbols, and the associated `labels` for each state.
+
+```python
+from dfax import DFAx
+
+dfax = DFAx(
+	start=0, # State referred to as 0 is the initial state
+	transitions=jnp.array([
+		[1, 2, 0, 0, 0],
+		[1, 1, 1, 1, 1],
+		[2, 2, 2, 2, 2],
+	]), # Max number of states is 3 and number of tokens is 5
+	labels=jnp.array([False, True, False]) # State labels
+) # Returns a DFA
+```
+
+Take transitions on the DFA using a given symbol.
+
+```python
+dfax = dfax.advance(0) # Returns the resulting DFA after reading the symbol referred to as 0
+```
+
+Minimize DFAs.
+
+```python
+dfax = dfax.minimize() # Returns a canonical minimal DFA
+```
+
+
+Canonicalize DFAs by relabeling states based on a BFS search.
+
+```python
+dfax = dfax.canonicalize() # Returns a canonical DFA
+```
+
+Mutate DFAs by randomly toggling entries in the transition matrix.
+
+```python
+import jax
+
+key = jax.random.PRNGKey(0)
+dfax = dfax.mutate(key) # Returns a mutated DFA
+```
+
+Perform syntactic equality check between DFAs.
+
+```python
+dfax1 == dfax2
+```
+
+Perform semantic equality check between DFAs.
+
+```python
+dfax1.minimize() == dfax2.minimize()
+```
+
+Use DFAs as reward functions. With ternary semantics, reward is (i) `+1` if the `start` state has label `True`, (ii) `-1` if the `start` state has label `False` and is a sink state, and (iii) `0` otherwise. With binary semantics, `0` is returned instead of `-1`.
+
+```python
+dfax.reward() # Returns a ternary reward
+dfax.reward(binary=True) # Returns a binary reward
+```
+
+
+Sample from different DFA distributions: `Reach` samples DFAs ordering alphabet symbols, `ReachAvoid` samples `Reach` DFAs but also includes `Avoid` constraints, and `ReachAvoidDerived` samples randomly mutated `Reach` and `ReachAvoid` DFAs.
+
+```python
+import jax
+from dfax.samplers import ReachSampler, ReachAvoidSampler, RADSampler
+
+key = jax.random.PRNGKey(0)
+sampler = ReachAvoidSampler()
+
+dfax = sampler.sample(key)
+```
+
+
+Define your own DFA samplers by overloading `DFASampler `.
+
+```python
+@struct.dataclass
+class MySampler(DFASampler):
+    @partial(jax.jit, static_argnums=(0,))
+    def sample(self, key: chex.PRNGKey) -> DFAx:
+        # Write sampling code and return sampled DFA
+```
+
+Visualize DFAs.
+
+```python
+from dfax.utils import visualize
+visualize(dfax)
+```
+
+
+This project is a JAX extension of [dfa](https://github.com/mvcisback/dfa). Therefore, we include helper methods for translating `DFAx` objects to and from `DFA` objects.
+
+```python
+from dfax import dfa2dfax, dfax2dfa
+
+dfa = dfax2dfa(dfax) # Create DFA from DFAx
+dfax = dfa2dfax(dfa) # Create DFAx from DFA
+```
+
+## In progress
+
+Currently, we are working on implementing Boolean operations on `DFAx` objects, e.g., conjunction, disjunction, etc. If there are other functionalities you would like to have in this package, create pull request or contact us to work together!
+
+
+
+
+

dfax-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+dfax/__init__.py,sha256=Z7ML6_mY6woYNNoO8G-dqRwj9iLITiiArFP_PFNhmVY,48
+dfax/dfax.py,sha256=dwBrsjSgrhBikvERw8aPL1fqTfQajr3YNOCTe1rsDFw,12486
+dfax/samplers.py,sha256=r9b48rioT54T6rmzakDLmzbKUZJ59JtwPjQutsW5vxg,7263
+dfax/utils.py,sha256=tEoruuxANe6ZIuNC7cqgTRef10VNbva7HNbgIXJGdGo,6732
+dfax-0.1.0.dist-info/METADATA,sha256=wzgpWzITxHhyy4Okwe7Jp3CH9WtW8SpnsOpabCA3cbk,3523
+dfax-0.1.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+dfax-0.1.0.dist-info/licenses/LICENSE,sha256=Cvu0BZqt3rcFFv70hcFDgD_y8ryOKW85F-qGRfYI4iM,1071
+dfax-0.1.0.dist-info/RECORD,,

dfax-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.28.0
+Root-Is-Purelib: true
+Tag: py3-none-any

dfax-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 RAD-Embeddings
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.