PINN for SWE

Start here

What should come after this?

• A token is one item from a fixed list.
• Our toy model must pick 0 or 1 or 2 here.
• A closing bracket is not allowed yet.

print(

Trace or excerpt. Download the full tested files to run them.

Start here

What rules add to training

• Data shows examples of valid programs.
• Rules tell us which choices are invalid.
• We can use both signals to train one model.

Start here

This token lab is not a physics PINN

• Grammar rules constrain discrete token choices.
• A physics PINN checks a differential equation.
• Both add known rules to learning. Their rules are not the same.

token lab: data loss + grammar loss
real PINN: equation loss + start loss

Trace or excerpt. Download the full tested files to run them.

Start here

A parser can already enforce grammar

• We do not need a neural net to check these fixed rules.
• The model learns choices among valid next tokens.
• The small language makes the training code easy to inspect.

print(0)      # one valid choice
print(1 + 2)  # another valid choice

Trace or excerpt. Download the full tested files to run them.

Start here

The route through this lesson

• Start with tokens and generated programs.
• Train a model with data and rules.
• Compare training rules with an inference mask.
• Move from grammar rules to a time equation.

Tokens and data

Our language is a tiny Python subset

• Each line prints one number or one two-number sum.
• Input numbers are only 0 through 2.
• We leave out all other Python features.

print(0)
print(1 + 2)
print(2 + 2)

Trace or excerpt. Download the full tested files to run them.

Tokens and data

A subset is not all of Python

• Full Python allows this code.
• Our tiny grammar does not.
• Failing this toy rule does not prove real Python code is invalid.

print()
print(4)
print("hello")

Trace or excerpt. Download the full tested files to run them.

Tokens and data

Text becomes a token list

• Token (one item from a fixed list).
• The generator makes token lists directly.
• It does not tokenize arbitrary Python files.

text:   print(1 + 2)
tokens: print  (  1  +  2  )  NEWLINE

Trace or excerpt. Download the full tested files to run them.

Tokens and data

The model has 11 tokens

• Vocabulary (all tokens this model can emit).
• Eight tokens represent code.
• Three special tokens help training and generation.

code = ["print", "(", "0", "1",
        "2", "+", ")", "\n"]
special = ["<PAD>", "<BOS>", "<EOS>"]

Trace or excerpt. Download the full tested files to run them.

Tokens and data

Special tokens are not Python source

• BOS marks the start of a program.
• EOS marks the end of a program.
• PAD fills unused space in a batch.
• We remove these markers before showing Python source.

<BOS> print ( 1 ) NEWLINE <EOS>

Trace or excerpt. Download the full tested files to run them.

Tokens and data

A token ID is only a lookup key

• Each token gets one integer ID.
• ID 5 does not mean a physical value of five.
• The model uses that ID to find a learned vector.

TOKEN_ID["0"] = 5
TOKEN_ID["1"] = 6
TOKEN_ID["2"] = 7

Trace or excerpt. Download the full tested files to run them.

Tokens and data

A prefix is the part already chosen

• Prefix (tokens that have already been chosen).
• The model can read the prefix.
• It cannot read the token it must predict next.

prefix: <BOS> print (
target: 1

Trace or excerpt. Download the full tested files to run them.

Tokens and data

One program gives several training steps

• We predict after every input position.
• Inputs and targets are shifted by one token.
• The target is never copied into its own input prefix.

inputs  = [BOS, print, (, 1, ), NL]
targets = [print, (, 1, ), NL, EOS]

Trace or excerpt. Download the full tested files to run them.

Tokens and data

The same prefix can have many valid targets

• All three continuations are valid.
• The data picks one target for each example.
• The model learns a distribution, not a single grammar answer.

print(0)
print(1)
print(2)

Trace or excerpt. Download the full tested files to run them.

Tokens and data

We make all data with code

• There are 12 possible statements in this subset.
• Programs have one to three statements.
• We shuffle 1,884 possible programs and keep 192.

3 single numbers + 9 sums = 12 lines
12 + 12**2 + 12**3 = 1884 programs

Trace or excerpt. Download the full tested files to run them.

Tokens and data

Split by whole program

• 96 programs train the model.
• 32 programs are available for validation.
• 64 programs form the held-out test split.
• Short prefixes can repeat across splits. Whole programs cannot.

Tokens and data

Inspect the generated JSONL data

• JSONL stores one JSON object on each line.
• Each row stores its split and token list. It also stores Python source.
• The loader checks grammar and duplicate programs.

{"split": "train",
 "tokens": ["<BOS>", "print", "(",
            "1", ")", "\n", "<EOS>"],
 "code": "print ( 1 )\n"}

Trace or excerpt. Download the full tested files to run them.

Tokens and data

Padding must not count as training data

• Long and short programs share a batch.
• PAD fills the shorter rows.
• Both loss terms ignore PAD targets.

active = targets != PAD
ce = cross_entropy(
    logits[active], targets[active]
)

Trace or excerpt. Download the full tested files to run them.

The network

A neural net is a function with learned weights

• We supply token IDs.
• The network supplies a score for each next token.
• Training changes weights to improve those scores.

The network

Embedding turns IDs into learned numbers

• Embedding (learned numbers that stand for tokens).
• Each token has an eight-number vector.
• Those vectors are learned during training.

embedding = nn.Embedding(11, 8)
vectors = embedding(token_ids)

Trace or excerpt. Download the full tested files to run them.

The network

The GRU reads tokens from left to right

• GRU (small network that tracks prior tokens).
• It updates its context after each token.
• No future token enters an earlier prediction.

rnn = nn.GRU(8, 16, batch_first=True)
hidden, _ = rnn(vectors)

Trace or excerpt. Download the full tested files to run them.

The network

Hidden state stores the prefix context

• Hidden state (numbers that store the prefix context).
• Our state has 16 learned numbers.
• These numbers are not the parser's named states.

parser state: first number needed
GRU state:    16 learned values

Trace or excerpt. Download the full tested files to run them.

The network

The output layer makes 11 scores

• Logit (raw score before turning into probability).
• A linear layer maps 16 values to 11 logits.
• A high logit means a stronger token preference.

head = nn.Linear(16, 11)
logits = head(hidden)

Trace or excerpt. Download the full tested files to run them.

The network

The complete model is small

• This model has 1,523 trainable values.
• Its input shape is batch by time.
• Its output adds one score per vocabulary token.

ids:    [batch, time]
logits: [batch, time, 11]

Trace or excerpt. Download the full tested files to run them.

The network

Softmax turns scores into probabilities

• Softmax (turns raw scores into token chances).
• Every probability is between zero and one.
• The probabilities sum to one at each position.

probabilities = logits.softmax(dim=-1)

Trace or excerpt. Download the full tested files to run them.

The network

Inspect real probabilities from our models

• Choose a measured prefix and model.
• Green tokens are allowed by the parser.
• The values come from the actual trained network.

Data loss

Loss gives training one score to reduce

• Loss (number that scores a model's error).
• A lower score is better for that objective.
• Different loss terms check different failures.

total = data_loss + weight * rule_loss

Trace or excerpt. Download the full tested files to run them.

Data loss

Cross-entropy rewards the known target

• Cross-entropy (loss for a known next token).
• It penalizes a low chance for the target token.
• A valid but different token still differs from this example's label.

target = "1"
loss = -log(probability_of_1)

Trace or excerpt. Download the full tested files to run them.

Data loss

A numerical data-loss example

• The example target is token 1.
• The model assigns it probability 0.2.
• The cross-entropy is about 1.609.

data_loss = -log(0.2)
          = 1.609

Trace or excerpt. Download the full tested files to run them.

Data loss

Data-only training does not enforce grammar

• The model can give invalid tokens nonzero probability.
• Cross-entropy discourages wrong targets indirectly.
• A grammar loss makes that penalty explicit.

prefix: print(
allowed: 0, 1, 2
invalid: ), +, EOS, ...

Trace or excerpt. Download the full tested files to run them.

Grammar loss

Grammar defines valid token sequences

• Grammar (rules that define valid token sequences).
• Our parser tracks which part of a statement comes next.
• It uses the known prefix, not the target label.

Grammar loss

After the opening bracket

• We require the first number.
• Only the number tokens 0 through 2 are legal.
• This is a rule of our subset, not a rule for all Python.

allowed_next(["<BOS>", "print", "("])
# {"0", "1", "2"}

Trace or excerpt. Download the full tested files to run them.

Grammar loss

After the first number

• A plus starts a two-number sum.
• A closing bracket ends a single-number expression.
• The prefix has two legal next-token choices.

prefix = ["<BOS>", "print", "(", "1"]
allowed_next(prefix)
# {"+", ")"}

Trace or excerpt. Download the full tested files to run them.

Grammar loss

After a complete line

• The next token can begin another print line.
• EOS can end the program.
• The grammar allows more lines than the dataset generator makes.

print(1) NEWLINE
next: print or EOS

Trace or excerpt. Download the full tested files to run them.

Grammar loss

Make a legal-token mask

• The mask has one boolean per vocabulary token.
• True means this token is allowed now.
• This mask does not enter the neural network's input.

legal = [token in allowed
         for token in VOCAB]

Trace or excerpt. Download the full tested files to run them.

Grammar loss

Sum probability on invalid tokens

• Add the chances assigned to forbidden choices.
• That sum is the grammar loss at this position.
• It does not say which valid token should win.

invalid_mass = (
    probabilities * (~legal)
).sum(dim=-1)

Trace or excerpt. Download the full tested files to run them.

Grammar loss

A numerical grammar-loss example

• Allowed tokens have total probability 0.6.
• The forbidden bracket has probability 0.4.
• The grammar loss is 0.4.

P(0)=0.3  P(1)=0.2  P(2)=0.1
P(close_bracket)=0.4
rule_loss = 0.4

Trace or excerpt. Download the full tested files to run them.

Grammar loss

Combine the two losses

• Data loss points toward the example's target.
• Grammar loss points away from invalid choices.
• The weight controls the extra rule penalty.

total = ce + lambda_rule * invalid_mass

Trace or excerpt. Download the full tested files to run them.

Grammar loss

Change the weight in a small arithmetic example

• These values explain the formula.
• The slider changes the displayed loss only.
• It does not retrain the network.

Grammar loss

A boolean rule can still guide gradients

• The legal mask is fixed for this prefix.
• Token probabilities depend on the weights.
• The loss changes smoothly when those weights change.

mask: fixed booleans
probabilities: learned values
loss: differentiable in model weights

Trace or excerpt. Download the full tested files to run them.

Grammar loss

The rule is not a physical derivative

• We do not differentiate with respect to token ID 5.
• We differentiate the loss with respect to model weights.
• A PINN also differentiates a prediction with respect to a physical coordinate.

token model: d(loss) / d(weights)
PINN:        du / dt, then gradients
             with respect to weights

Trace or excerpt. Download the full tested files to run them.

Training

One training step has three phases

• Predict token scores.
• Compute the data and rule losses.
• Use gradients to update weights.

Training

Clear old weight gradients

• PyTorch accumulates gradients by default.
• This example wants a fresh gradient each step.
• Clear it before the forward pass.

optimizer.zero_grad(set_to_none=True)

Trace or excerpt. Download the full tested files to run them.

Training

Run the forward pass

• Input IDs contain known past tokens.
• The network returns one score vector per position.
• The parser supplies the separate loss mask.

logits = model(input_ids)
ce, rule = loss_terms(
    logits, targets, legal
)

Trace or excerpt. Download the full tested files to run them.

Training

Backprop finds weight gradients

• Gradient (change in loss when weights change).
• Backprop (finds weight gradients from the loss).
• Both loss terms contribute to the gradient.

total = ce + 4.0 * rule
total.backward()

Trace or excerpt. Download the full tested files to run them.

Training

The optimizer updates weights

• Optimizer (uses gradients to update model weights).
• Adam is the optimizer in this example.
• An update is not a proof that training will converge.

optimizer.step()

Trace or excerpt. Download the full tested files to run them.

Training

Teacher forcing uses known past tokens

• Teacher forcing (feeds known past tokens during training).
• Every prediction sees the real prefix from the dataset.
• It does not see future targets.

known prefix -> predict next token
next step uses the dataset's next token

Trace or excerpt. Download the full tested files to run them.

Training

Compare paired training runs fairly

• Both models start from the same weights.
• Both see the same programs and batch order.
• Only the grammar-loss weight changes.

data_only:   lambda_rule = 0
rule_guided: lambda_rule = 4

Trace or excerpt. Download the full tested files to run them.

Training

Held-out results from this run

• The displayed results come from seed 7 and 80 epochs.
• Test programs did not update the weights.
• One run does not prove rules always improve a model.

Training

Why next-token accuracy is not 100 percent

• Numbers and expression forms are randomly chosen in the data.
• A prefix may have several legal continuations.
• Predicting a different legal token can count as a label miss.

prefix: print(
true label: 1
prediction: 2
syntax valid, label miss

Trace or excerpt. Download the full tested files to run them.

Generation

Generation uses the model's own choices

• Autoregressive (uses its own past token choices).
• Append the sampled token to the prefix.
• Then predict again until EOS or a token budget.

Generation

Soft training rules are not hard guarantees

• A forbidden token can still have nonzero probability.
• Sampling can pick that token.
• We must test complete programs, not just known-prefix predictions.

low invalid probability != zero risk

Trace or excerpt. Download the full tested files to run them.

Generation

A hard mask blocks forbidden tokens now

• Hard mask (sets forbidden token chances to zero).
• Replace invalid logits with negative infinity.
• Then sample from the allowed set.

logits = logits.masked_fill(
    ~legal, -torch.inf
)
p = logits.softmax(-1)

Trace or excerpt. Download the full tested files to run them.

Generation

Hard masking is an inference rule

• It does not change the model weights.
• It can be used with either trained model.
• It is different from a rule loss used during training.

training penalty: changes weights
inference mask:   changes this choice

Trace or excerpt. Download the full tested files to run them.

Generation

Legal choices do not guarantee completion

• The parser can keep allowing another line.
• The model may fail to emit EOS within 32 tokens.
• We report a partial program as incomplete.

stop_reason = "token_budget"
complete = False
code = None

Trace or excerpt. Download the full tested files to run them.

Generation

Full-program results reveal another limit

• We sample 100 programs with each decoding method.
• The report separates invalid choices from budget stops.
• A hard mask protects syntax transitions, not termination.

Generation

Syntax and meaning are different tests

• The grammar accepts this program.
• Its output is 4.
• Syntax does not prove the program solves a user's task.

print(2 + 2)
# output: 4

Trace or excerpt. Download the full tested files to run them.

Real PINN

Now replace grammar with a time equation

• PINN (network trained with a physics equation).
• The input is a continuous time value.
• The output is the amount left in a decay process.

Real PINN

Use a normalized decay process

• The process starts at one unit.
• Its rate of change is minus its current amount.
• We use time from zero to two in normalized units.

u'(t) = -u(t)
u(0) = 1
0 <= t <= 2

Trace or excerpt. Download the full tested files to run them.

Real PINN

The network predicts a number, not a token

• It has one time input and one value output.
• Two hidden layers use Tanh for a smooth function.
• This model has 321 trainable values.

Linear(1, 16), Tanh()
Linear(16, 16), Tanh()
Linear(16, 1)

Trace or excerpt. Download the full tested files to run them.

Real PINN

A derivative measures change with time

• Derivative (how fast output changes with time).
• We ask for du/dt at each sample time.
• The coordinate t needs gradient tracking.

t = torch.rand(64, 1) * 2
t.requires_grad_()
u = model(t)

Trace or excerpt. Download the full tested files to run them.

Real PINN

Autograd computes du/dt

• Autograd (tracks operations to find their derivatives).
• The network treats rows as independent samples.
• That makes this gradient the per-row time derivative.

du_dt = torch.autograd.grad(
    u, t,
    grad_outputs=torch.ones_like(u),
    create_graph=True
)[0]

Trace or excerpt. Download the full tested files to run them.

Real PINN

Why create_graph is True

• We will put the time derivative inside the loss.
• Training must also differentiate that derivative with respect to weights.
• create_graph=True keeps that second gradient path.

time -> prediction -> time derivative
                    -> loss -> weights

Trace or excerpt. Download the full tested files to run them.

Real PINN

Make an equation residual

• Residual (amount that violates the stated equation).
• The equation is du/dt plus u equals zero.
• Square the residual and average it over time points.

residual = du_dt + u
physics_loss = residual.square().mean()

Trace or excerpt. Download the full tested files to run them.

Real PINN

The equation alone accepts the wrong zero curve

• A constant zero has zero derivative.
• It gives zero equation residual.
• It violates the required start value of one.

u(t) = 0
u'(t) + u(t) = 0
u(0) != 1

Trace or excerpt. Download the full tested files to run them.

Real PINN

Add the start condition

• Initial condition (known value at the starting time).
• We require u(0) to equal one.
• The start condition excludes the zero solution.

initial_loss = (
    model(torch.zeros(1, 1)) - 1
).square().mean()

Trace or excerpt. Download the full tested files to run them.

Real PINN

The start condition alone is not enough either

• A constant one satisfies the start value.
• It does not decay with time.
• Both checks are needed.

u(t) = 1
u(0) = 1
u'(t) + u(t) = 1

Trace or excerpt. Download the full tested files to run them.

Real PINN

Train with equation and start losses

• This run uses a start-loss weight of ten.
• Time points come from code, not sensors.
• There are no other time/value labels in training.

total = physics_loss + 10 * initial_loss
total.backward()
optimizer.step()

Trace or excerpt. Download the full tested files to run them.

Real PINN

Check at fresh times

• The exact reference is exp(-t).
• We use that curve only after training.
• A low training residual does not prove accuracy everywhere.

test_times = linspace(0, 2, 101)
reference = exp(-test_times)
error = prediction - reference

Trace or excerpt. Download the full tested files to run them.

Real PINN

The measured decay curve

• This run used seed 7 and 1,500 epochs.
• The chart compares predictions with the exact reference.
• The reference curve did not train the model.

Evaluation RMSE: 0.00221. Maximum absolute error: 0.00363.

Real PINN

What transfers from the token lab

• A model learns from a loss.
• Known rules add another loss term.
• The rule type decides how that extra term is built.

grammar: invalid probability mass
physics: squared equation residual

Trace or excerpt. Download the full tested files to run them.

Run the labs

Download the complete programs

• Both scripts run on CPU.
• The token script makes its own dataset.
• The decay script needs no sensor records.

Run the labs

Set up PyTorch

• Use Python 3.10 or newer.
• Create a fresh environment in your own folder.
• On Windows use .venv\Scripts\activate.

python3 -m venv .venv
source .venv/bin/activate
python -m pip install torch

Trace or excerpt. Download the full tested files to run them.

Run the labs

Train the token models

• Download token_rules.py into that folder.
• The command trains both paired models.
• It saves the data and model weights. It also saves results.

python token_rules.py train
# default output: runs/tokens/
# inspect runs/tokens/results.json

Trace or excerpt. Download the full tested files to run them.

Run the labs

Try a different rule weight

• Keep the seed and data fixed.
• Change only the grammar-loss weight.
• Compare probability mass and full-program behavior.

python token_rules.py train \
  --rule-weight 1 \
  --output-dir runs/weight_1

Trace or excerpt. Download the full tested files to run them.

Run the labs

Train the real PINN

• Download decay_pinn.py into the same folder.
• The model learns the decay law and start value.
• The saved CSV contains evaluation values, not training labels.

python decay_pinn.py
# inspect runs/decay/results.json
# inspect runs/decay/evaluation.csv

Trace or excerpt. Download the full tested files to run them.

Run the labs

Practice by changing one rule

• First explain the current allowed-next set.
• Then permit one new number in the grammar and vocabulary.
• Regenerate data and test masks before training.
• Do not change the held-out results to make a claim look better.

exercise: add token "3"
update VOCAB, NUMBERS, and tests

Trace or excerpt. Download the full tested files to run them.

Summary

▶ Summary

• The token lab learns choices with a separate grammar penalty.
• A hard decoding mask can block invalid choices without retraining.
• A real PINN checks a continuous equation and its conditions.

Action items

▶ Action Items

1. Run the token lab and inspect invalid_mass in results.json.
2. Then run the decay lab and inspect its error.

State: two runnable labs are ready.
Next: train the token model.

Balbir Singh · 22 October 2026
Generated with the visualize skill

═══════════════════════════════════════