Customizing a TensorFlow operation

At a glance

Project structure

Source bundle/
├── hash_encoder/
│   ├── hash_encoder_kernel.cc
│   ├── hash_encoder_kernel.metal
│   ├── hash_encoder.py
│   └── mtl_hash_encoder_kernel.cc
├── tiny_nerf_hash.py
├── tiny_nerf_mlp.py
└── render_utils.py

Structure observations

• The entry/composition boundary and renderer or operation boundary are separate in the source; C++ TensorFlow op is the principal feature coordinator.
• GPU-specific logic stays in Metal shader files; shared headers bridge host/shader layouts where present.
• The tree above is intentionally pruned to composition, resource, and shader files.

Overall architecture

flowchart LR
    N1["Python layer/model"]
    N2["C++ TensorFlow op"]
    N3["Python layer and registered C++ custom operation"]
    N4["Metal hash-encoding kernel"]
    N1 --> N2
    N2 --> N3
    N3 --> N4

Reference code

hash_encoder/hash_encoder_kernel.cc:13 — feature handoff or setup anchor

REGISTER_OP("HashEncode")
    .Attr("T: {float}")
    // The input coordinates.
    .Input("inputs: float") 
    // The NGP embedding buffer.
    .Input("embeddings: T") 
    // The hashmap offsets of all levels.
    .Input("hashmap_offsets: int32")

Interpretation

The Python layer crosses TensorFlow’s custom-op ABI. C++ owns operation/kernel registration and the native dispatch path; the .metal file owns the parallel forward and backward algorithms.

Ownership and state

classDiagram
    NGP *-- HashEncoder : enc

Ownership evidence

tiny_nerf_hash.py:313 — NGP creates and stores enc

        self.enc = HashEncoder(n_dim=ngp_n_dim, n_levels=ngp_levels, n_feature=ngp_feature,
                               resolution_coarsest=32, log2_hashmap_size=19, resolution_finest=256)

        # The output feature length.
        ngp_output_channel = ngp_feature * ngp_levels

        # The geometry feature vector length for the color net.
        geometry_feat = 3

Class and protocol design

No local protocol abstraction is present; the sample uses concrete framework, operation, or model types at this boundary.

Access control

This sample does not use Swift private, fileprivate, or public for its native boundary. Objective-C/Objective-C++ use header versus implementation placement, C++ uses access specifiers/linkage, Python uses module conventions, and Metal entry points cross a compiled-library boundary; these are not Swift access levels.

Logic ownership and placement

Shader boundary reference

hash_encoder/hash_encoder_kernel.metal:184 — representative GPU entry/helper

kernel void HashEncodeForward(
    // ...
    uint3 tid [[thread_position_in_grid]]) {
    // ...
}

Design patterns

Naming conventions

• Role suffixes make ownership visible: HashEncoder.
• Method names describe setup or encoding actions: forward, forward_with_tensors_only, grad, __init__, call, OK, queue.
• Shader entry points use stage/operation names: HashEncodeForward, HashEncodeBackward.
• Names favor concrete domain roles and target-local types; no broad public library namespace is introduced.

Architecture takeaways

• Keep Python layer/model focused on composition; C++ TensorFlow op is the owner of Python layer and registered C++ custom operation.
• Treat Metal hash-encoding kernel as a separate execution/compilation boundary with explicit resource and data-layout contracts.
• The verified ownership edge is NGP → enc; broader exclusive ownership is not inferred.
• Access-control rationale follows concrete language boundaries rather than translating every header or shader symbol into Swift terms.

Source map

• Official Apple documentation