Customizing a TensorFlow operation

At a glance

Project structure

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

Structure observations

• Architecturally prominent files are ranked from entry points and role-named declarations; resource-only paths are omitted.
• Primary languages: C++, Metal, Python.
• The verified tree contains 0 project/configuration file(s) and 10 source declaration(s).

Overall architecture

flowchart LR
    N1["_hash_encode"]
    N2["TrainMonitor"]
    N3["Metal APIs"]
    N1 --> N2
    N2 --> N3

Reference code

hash_encoder/hash_encoder.py:41 — architecture anchor

class _hash_encode:
    @staticmethod
    def forward(inputs, embeddings, hashmap_offsets, level_scale_ratio, resolution_coarsest):

        log2_per_level_scale = np.log2(level_scale_ratio)

        # A custom gradient tensorflow function to bind the forward/backward kernel.
        @tf.custom_gradient
        def forward_with_tensors_only(inputs, embeddings, hashmap_offsets):
            # Forward kernel call.
            outputs = _backend.hash_encode(
                inputs, embeddings, hashmap_offsets, log2_per_level_scale, resolution_coarsest)

            def grad(incoming_gradients):
                # The shape of "incoming_gradients": [B, L * C]

                # Backward kernel call.
                grad_embeddings = _backend.hash_encode_grad(
                    incoming_gradients, inputs, embeddings, hashmap_offsets, log2_per_level_scale, resolution_coarsest)
                return None, grad_embeddings, None

            return outputs, grad

        return forward_with_tensors_only(inputs, embeddings, hashmap_offsets)

Interpretation

The arrows summarize the source-visible entry, role-named types or folders, and framework direction; when nodes come from structural folders, the sequence is a high-level interpretation rather than proof that every adjacent node calls the next. Ownership is claimed only where the next section cites a stored property or assignment. The diagram is intentionally limited to the dominant path into Metal.

Ownership and state

classDiagram
    _hash_encode --> MetalAPIs : uses

Ownership evidence

hash_encoder/hash_encoder_kernel.cc:11 — stored dependency or nearest verified ownership anchor

using namespace tensorflow;

Composition arrows indicate a source-visible construction expression or locally owned value state; aggregation means the owner stores or receives a dependency without proving exclusive lifetime ownership.

Concurrency, scheduling, and thread safety

Evidence limit: actor isolation, async/await, or Task creation does not by itself prove background-thread execution; Sendable conformance alone does not prove thread-safe mutation.

No source-visible execution, scheduling, or synchronization boundary was found in the indexed source.

@MainActor/MainActor.run, DispatchQueue.main, and RunLoop.main are reported as distinct isolation, queue, and event-loop mechanisms. A plain Task is kept separate from Task.detached; neither is labeled as a background thread.

State propagation, frameworks, and dependencies

Evidence limit: an import proves a source-level compilation dependency at the cited line; it does not prove runtime use, architectural adoption, or whether a Swift package is a direct application dependency.

receive(on:) describes downstream delivery scheduling, while subscribe(on:) describes upstream subscription/request/cancel scheduling. An import Combine alone establishes neither behavior nor a Store, reducer, Redux, or other application architecture.

Class and protocol design

tiny_nerf_hash.py:435 — representative type boundary

class TrainMonitor(keras.callbacks.Callback):
    # ...

No local protocol conformance is claimed as protocol-oriented design; external framework conformances are listed only as dependencies.

Access control

Reference code

hash_encoder/mtl_hash_encoder_kernel.cc:19 — representative boundary

@protocol TF_MetalStream

Swift declarations without a modifier are internal; explicit private, fileprivate, private(set), public, or open entries above are interpreted by language semantics. Objective-C/C samples instead rely on header and implementation boundaries, which are not equivalent to Swift lexical privacy.

Logic ownership and placement

Design patterns

Naming conventions

• Types: Monitor: TrainMonitor.
• Protocols: no local protocol declaration in the scanned source.
• Methods: forward, forward_with_tensors_only, grad, __init__, call, queue, currentCommandBuffer, commit.
• Files: feature/project roles rather than a strict one-type-per-file rule.

Architecture takeaways

• _hash_encode is the main source-visible entry or composition anchor for this sample.
• Framework work reaches tensorflow, Metal, dispatch, filesystem through a deliberately small high-level chain; the detailed API graph remains inside the cited implementation files.
• Stored-property evidence identifies lifecycle collaboration; it does not by itself prove exclusive object ownership.
• Access-control conclusions separate verified language visibility from the likely design rationale.
• The source does not justify labeling the design protocol-oriented.

Source map

• Official Apple documentation