```python
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.cpp_extension import load_inline

# Fused Triplet Margin Loss CUDA kernels.
#
# The whole loss is computed in one streaming pass over the inputs:
#   dist_pos = ||anchor - positive + eps||_2   (fused sub/square/reduce per row)
#   dist_neg = ||anchor - negative + eps||_2
#   loss_i   = max(dist_pos - dist_neg + margin, 0)
#   output   = mean(loss_i)                    (tiny second reduction kernel)
#
# Improvements over the previous version:
#   * A fully-unrolled fixed-trip-count kernel is dispatched when the
#     vectorized row length is an exact multiple of the block size
#     (covers the benchmark shape dim=8192), maximizing memory-level
#     parallelism and letting the compiler batch all loads.
#   * Four independent accumulators per distance break FADD dependency
#     chains in every kernel.
#   * Streaming loads (__ldcs) for the read-once inputs.
#   * Vectorized (float4) multi-accumulator mean kernel.
#   * 16B-alignment checks guard the vectorized paths, with a scalar
#     fallback for arbitrary shapes/strides.
# A matching fused backward kernel keeps the module fully trainable, and a
# dedicated eval entry point skips materializing the per-row distances.

_triplet_cuda_src = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>
#include <vector>
#include <cstdint>

#define TML_BLOCK 256
#define TML_WARPS (TML_BLOCK / 32)

__device__ __forceinline__ void acc4(const float4 a, const float4 b,
                                     const float eps,
                                     float& s0, float& s1,
                                     float& s2, float& s3) {
    float d;
    d = a.x - b.x + eps; s0 += d * d;
    d = a.y - b.y + eps; s1 += d * d;
    d = a.z - b.z + eps; s2 += d * d;
    d = a.w - b.w + eps; s3 += d * d;
}

template<bool SAVE_DISTS>
__device__ __forceinline__ void reduce_and_store(
    float sp, float sn, const int row,
    float* __restrict__ losses,
    float* __restrict__ dist_pos,
    float* __restrict__ dist_neg,
    const float margin)
{
    __shared__ float sh_p[TML_WARPS];
    __shared__ float sh_n[TML_WARPS];
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sp += __shfl_down_sync(0xffffffffu, sp, o);
        sn += __shfl_down_sync(0xffffffffu, sn, o);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) { sh_p[warp] = sp; sh_n[warp] = sn; }
    __syncthreads();
    if (warp == 0) {
        sp = (lane < TML_WARPS) ? sh_p[lane] : 0.0f;
        sn = (lane < TML_WARPS) ? sh_n[lane] : 0.0f;
        #pragma unroll
        for (int o = TML_WARPS / 2; o > 0; o >>= 1) {
            sp += __shfl_down_sync(0xffffffffu, sp, o);
            sn += __shfl_down_sync(0xffffffffu, sn, o);
        }
        if (lane == 0) {
            const float dp = sqrtf(sp);
            const float dn = sqrtf(sn);
            if (SAVE_DISTS) {
                dist_pos[row] = dp;
                dist_neg[row] = dn;
            }
            losses[row] = fmaxf(dp - dn + margin, 0.0f);
        }
    }
}

// Fully-unrolled kernel: row vector length == VT * TML_BLOCK (in float4s).
template<bool SAVE_DISTS, int VT>
__global__ __launch_bounds__(TML_BLOCK) void tml_forward_fixed_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    float* __restrict__ dist_pos,
    float* __restrict__ dist_neg,
    const float margin,
    const float eps)
{
    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)(VT * TML_BLOCK);
    const float4* a4 = reinterpret_cast<const float4*>(anchor) + base;
    const float4* p4 = reinterpret_cast<const float4*>(positive) + base;
    const float4* n4 = reinterpret_cast<const float4*>(negative) + base;

    float sp0 = 0.0f, sp1 = 0.0f, sp2 = 0.0f, sp3 = 0.0f;
    float sn0 = 0.0f, sn1 = 0.0f, sn2 = 0.0f, sn3 = 0.0f;

    #pragma unroll
    for (int k = 0; k < VT; ++k) {
        const int i = threadIdx.x + k * TML_BLOCK;
        const float4 av = __ldcs(a4 + i);
        const float4 pv = __ldcs(p4 + i);
        const float4 nv = __ldcs(n4 + i);
        acc4(av, pv, eps, sp0, sp1, sp2, sp3);
        acc4(av, nv, eps, sn0, sn1, sn2, sn3);
    }
    const float sp = (sp0 + sp1) + (sp2 + sp3);
    const float sn = (sn0 + sn1) + (sn2 + sn3);
    reduce_and_store<SAVE_DISTS>(sp, sn, row, losses, dist_pos, dist_neg, margin);
}

// Generic float4 kernel (requires dim % 4 == 0 and 16B-aligned rows).
template<bool SAVE_DISTS>
__global__ __launch_bounds__(TML_BLOCK) void tml_forward_vec_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    float* __restrict__ dist_pos,
    float* __restrict__ dist_neg,
    const int dim,
    const float margin,
    const float eps)
{
    const int row = blockIdx.x;
    const int vdim = dim >> 2;
    const size_t base = (size_t)row * (size_t)vdim;
    const float4* a4 = reinterpret_cast<const float4*>(anchor) + base;
    const float4* p4 = reinterpret_cast<const float4*>(positive) + base;
    const float4* n4 = reinterpret_cast<const float4*>(negative) + base;

    float sp0 = 0.0f, sp1 = 0.0f, sp2 = 0.0f, sp3 = 0.0f;
    float sn0 = 0.0f, sn1 = 0.0f, sn2 = 0.0f, sn3 = 0.0f;

    int i = threadIdx.x;
    // Unrolled x4: keeps twelve independent 16-byte loads in flight.
    for (; i + 3 * TML_BLOCK < vdim; i += 4 * TML_BLOCK) {
        const float4 av0 = __ldcs(a4 + i);
        const float4 av1 = __ldcs(a4 + i + TML_BLOCK);
        const float4 av2 = __ldcs(a4 + i + 2 * TML_BLOCK);
        const float4 av3 = __ldcs(a4 + i + 3 * TML_BLOCK);
        const float4 pv0 = __ldcs(p4 + i);
        const float4 pv1 = __ldcs(p4 + i + TML_BLOCK);
        const float4 pv2 = __ldcs(p4 + i + 2 * TML_BLOCK);
        const float4 pv3 = __ldcs(p4 + i + 3 * TML_BLOCK);
        const float4 nv0 = __ldcs(n4 + i);
        const float4 nv1 = __ldcs(n4 + i + TML_BLOCK);
        const float4 nv2 = __ldcs(n4 + i + 2 * TML_BLOCK);
        const float4 nv3 = __ldcs(n4 + i + 3 * TML_BLOCK);
        acc4(av0, pv0, eps, sp0, sp1, sp2, sp3);
        acc4(av1, pv1, eps, sp0, sp1, sp2, sp3);
        acc4(av2, pv2, eps, sp0, sp1, sp2, sp3);
        acc4(av3, pv3, eps, sp0, sp1, sp2, sp3);
        acc4(av0, nv0, eps, sn0, sn1, sn2, sn3);
        acc4(av1, nv1, eps, sn0, sn1, sn2, sn3);
        acc4(av2, nv2, eps, sn0, sn1, sn2, sn3);
        acc4(av3, nv3, eps, sn0, sn1, sn2, sn3);
    }
    for (; i < vdim; i += TML_BLOCK) {
        const float4 av = __ldcs(a4 + i);
        const float4 pv = __ldcs(p4 + i);
        const float4 nv = __ldcs(n4 + i);
        acc4(av, pv, eps, sp0, sp1, sp2, sp3);
        acc4(av, nv, eps, sn0, sn1, sn2, sn3);
    }
    const float sp = (sp0 + sp1) + (sp2 + sp3);
    const float sn = (sn0 + sn1) + (sn2 + sn3);
    reduce_and_store<SAVE_DISTS>(sp, sn, row, losses, dist_pos, dist_neg, margin);
}

// Scalar fallback for arbitrary dim / alignment.
template<bool SAVE_DISTS>
__global__ __launch_bounds__(TML_BLOCK) void tml_forward_scalar_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    float* __restrict__ losses,
    float* __restrict__ dist_pos,
    float* __restrict__ dist_neg,
    const int dim,
    const float margin,
    const float eps)
{
    const int row = blockIdx.x;
    const size_t base = (size_t)row * (size_t)dim;
    const float* a = anchor + base;
    const float* p = positive + base;
    const float* n = negative + base;

    float sp0 = 0.0f, sp1 = 0.0f;
    float sn0 = 0.0f, sn1 = 0.0f;
    int i = threadIdx.x;
    for (; i + TML_BLOCK < dim; i += 2 * TML_BLOCK) {
        const float av0 = __ldg(a + i);
        const float av1 = __ldg(a + i + TML_BLOCK);
        const float dp0 = av0 - __ldg(p + i) + eps;
        const float dp1 = av1 - __ldg(p + i + TML_BLOCK) + eps;
        const float dn0 = av0 - __ldg(n + i) + eps;
        const float dn1 = av1 - __ldg(n + i + TML_BLOCK) + eps;
        sp0 += dp0 * dp0; sp1 += dp1 * dp1;
        sn0 += dn0 * dn0; sn1 += dn1 * dn1;
    }
    for (; i < dim; i += TML_BLOCK) {
        const float av = __ldg(a + i);
        const float dp = av - __ldg(p + i) + eps;
        const float dn = av - __ldg(n + i) + eps;
        sp0 += dp * dp;
        sn0 += dn * dn;
    }
    reduce_and_store<SAVE_DISTS>(sp0 + sp1, sn0 + sn1, row,
                                 losses, dist_pos, dist_neg, margin);
}

__global__ __launch_bounds__(TML_BLOCK) void tml_mean_kernel(
    const float* __restrict__ losses,
    float* __restrict__ out,
    const int batch)
{
    // Accumulate in double for an accurate mean; float4 loads.
    double s0 = 0.0, s1 = 0.0, s2 = 0.0, s3 = 0.0;
    const int vbatch = batch >> 2;
    const float4* l4 = reinterpret_cast<const float4*>(losses);
    for (int i = threadIdx.x; i < vbatch; i += TML_BLOCK) {
        const float4 v = __ldg(l4 + i);
        s0 += (double)v.x;
        s1 += (double)v.y;
        s2 += (double)v.z;
        s3 += (double)v.w;
    }
    double sum = (s0 + s1) + (s2 + s3);
    for (int i = (vbatch << 2) + threadIdx.x; i < batch; i += TML_BLOCK) {
        sum += (double)__ldg(losses + i);
    }
    __shared__ double sh[TML_WARPS];
    #pragma unroll
    for (int o = 16; o > 0; o >>= 1) {
        sum += __shfl_down_sync(0xffffffffu, sum, o);
    }
    const int lane = threadIdx.x & 31;
    const int warp = threadIdx.x >> 5;
    if (lane == 0) sh[warp] = sum;
    __syncthreads();
    if (warp == 0) {
        sum = (lane < TML_WARPS) ? sh[lane] : 0.0;
        #pragma unroll
        for (int o = TML_WARPS / 2; o > 0; o >>= 1) {
            sum += __shfl_down_sync(0xffffffffu, sum, o);
        }
        if (lane == 0) out[0] = (float)(sum / (double)batch);
    }
}

__global__ __launch_bounds__(TML_BLOCK) void tml_backward_kernel(
    const float* __restrict__ anchor,
    const float* __restrict__ positive,
    const float* __restrict__ negative,
    const float* __restrict__ dist_pos,
    const float* __restrict__ dist_neg,
    const float* __restrict__ grad_out,
    float* __restrict__ grad_anchor,
    float* __restrict__ grad_positive,
    float* __restrict__ grad_negative,
    const int batch,
    const int dim,
    const float margin,
    const float eps)
{
    const int row = blockIdx.x;
    const float dp = __ldg(dist_pos + row);
    const float dn = __ldg(dist_neg + row);
    const float scale = __ldg(grad_out) / (float)batch;
    const float act = (dp - dn + margin > 0.0f) ? scale : 0.0f;
    const float cp = (dp > 0.0f) ? (act / dp) : 0.0f;
    const float cn = (dn > 0.0f) ? (act / dn) : 0.0f;

    const size_t base = (size_t)row * (size_t)dim;
    const float* a = anchor + base;
    const float* p = positive + base;
    const float* n = negative + base;
    float* ga = grad_anchor + base;
    float* gp = grad_positive + base;
    float* gn = grad_negative + base;

    const bool aligned =
        ((((uintptr_t)a) | ((uintptr_t)p) | ((uintptr_t)n) |
          ((uintptr_t)ga) | ((uintptr_t)gp) | ((uintptr_t)gn)) & (uintptr_t)15) == 0;

    if (((dim & 3) == 0) && aligned) {
        const int vdim = dim >> 2;
        const float4* a4 = reinterpret_cast<const float4*>(a);
        const float4* p4 = reinterpret_cast<const float4*>(p);
        const float4* n4 = reinterpret_cast<const float4*>(n);
        float4* ga4 = reinterpret_cast<float4*>(ga);
        float4* gp4 = reinterpret_cast<float4*>(gp);
        float4* gn4 = reinterpret_cast<float4*>(gn);
        for (int i = threadIdx.x; i < vdim; i += TML_BLOCK) {
            const float4 av = __ldg(a4 + i);
            const float4 pv = __ldg(p4 + i);
            const float4 nv = __ldg(n4 + i);
            float4 gav, gpv, gnv;
            float dap, dan;
            dap = av.x - pv.x + eps; dan = av.x - nv.x + eps;
            gav.x = cp * dap - cn * dan; gpv.x = -cp * dap; gnv.x = cn * dan;
            dap = av.y - pv.y + eps; dan = av.y - nv.y + eps;
            gav.y = cp * dap - cn * dan; gpv.y = -cp * dap; gnv.y = cn * dan;
            dap = av.z - pv.z + eps; dan = av.z - nv.z + eps;
            gav.z = cp * dap - cn * dan; gpv.z = -cp * dap; gnv.z = cn * dan;
            dap = av.w - pv.w + eps; dan = av.w - nv.w + eps;
            gav.w = cp * dap - cn * dan; gpv.w = -cp * dap; gnv.w = cn * dan;
            ga4[i] = gav;
            gp4[i] = gpv;
            gn4[i] = gnv;
        }
    } else {
        for (int i = threadIdx.x; i < dim; i += TML_BLOCK) {
            const float av = __ldg(a + i);
            const float dap = av - __ldg(p + i) + eps;
            const float dan = av - __ldg(n + i) + eps;
            ga[i] = cp * dap - cn * dan;
            gp[i] = -cp * dap;
            gn[i] = cn * dan;
        }
    }
}

template<bool SAVE_DISTS>
static inline void launch_tml_forward(
    const float* a, const float* p, const float* n,
    float* losses, float* dpos, float* dneg,
    const int batch, const int dim,
    const float margin, const float eps,
    cudaStream_t stream)
{
    const bool aligned =
        ((((uintptr_t)a) | ((uintptr_t)p) | ((uintptr_t)n)) & (uintptr_t)15) == 0;
    if (aligned && ((dim & 3) == 0)) {
        const int vdim = dim >> 2;
        if ((vdim % TML_BLOCK) == 0) {
            const int vt = vdim / TML_BLOCK;
            switch (vt) {
            case 1:
                tml_forward_fixed_kernel<SAVE_DISTS, 1><<<batch, TML_BLOCK, 0, stream>>>(
                    a, p, n, losses, dpos, dneg, margin, eps);
                return;
            case 2:
                tml_forward_fixed_kernel<SAVE_DISTS, 2><<<batch, TML_BLOCK, 0, stream>>>(
                    a, p, n, losses, dpos, dneg, margin, eps);
                return;
            case 4:
                tml_forward_fixed_kernel<SAVE_DISTS, 4><<<batch, TML_BLOCK, 0, stream>>>(
                    a, p, n, losses, dpos, dneg, margin, eps);
                return;
            case 8:
                tml_forward_fixed_kernel<SAVE_DISTS, 8><<<batch, TML_BLOCK, 0, stream>>>(
                    a, p, n, losses, dpos, dneg, margin, eps);
                return;
            default:
                break;
            }
        }
        tml_forward_vec_kernel<SAVE_DISTS><<<batch, TML_BLOCK, 0, stream>>>(
            a, p, n, losses, dpos, dneg, dim, margin, eps);
    } else {
        tml_forward_scalar_kernel<SAVE_DISTS><<<batch, TML_BLOCK, 0, stream>>>(
            a, p, n, losses, dpos, dneg, dim, margin, eps);
    }
}

torch::Tensor triplet_forward_eval(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps)
{
    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() && negative.is_cuda(),
                "inputs must be CUDA tensors");
    TORCH_CHECK(anchor.scalar_type() == at::kFloat &&
                positive.scalar_type() == at::kFloat &&
                negative.scalar_type() == at::kFloat,
                "inputs must be float32");

    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();

    const int64_t batch = a.size(0);
    const int64_t dim = batch > 0 ? (a.numel() / batch) : 0;

    auto opts = a.options();
    auto losses = at::empty({batch}, opts);
    auto out    = at::empty({}, opts);

    if (batch == 0 || dim == 0) {
        out.fill_(0.0f);
        return out;
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    launch_tml_forward<false>(
        a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
        losses.data_ptr<float>(), nullptr, nullptr,
        (int)batch, (int)dim, (float)margin, (float)eps, stream);
    tml_mean_kernel<<<1, TML_BLOCK, 0, stream>>>(
        losses.data_ptr<float>(), out.data_ptr<float>(), (int)batch);

    return out;
}

std::vector<torch::Tensor> triplet_forward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps)
{
    TORCH_CHECK(anchor.is_cuda() && positive.is_cuda() && negative.is_cuda(),
                "inputs must be CUDA tensors");
    TORCH_CHECK(anchor.scalar_type() == at::kFloat &&
                positive.scalar_type() == at::kFloat &&
                negative.scalar_type() == at::kFloat,
                "inputs must be float32");

    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();

    const int64_t batch = a.size(0);
    const int64_t dim = batch > 0 ? (a.numel() / batch) : 0;

    auto opts = a.options();
    auto losses   = at::empty({batch}, opts);
    auto dist_pos = at::empty({batch}, opts);
    auto dist_neg = at::empty({batch}, opts);
    auto out      = at::empty({}, opts);

    if (batch == 0 || dim == 0) {
        out.fill_(0.0f);
        return {out, dist_pos, dist_neg};
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    launch_tml_forward<true>(
        a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
        losses.data_ptr<float>(), dist_pos.data_ptr<float>(), dist_neg.data_ptr<float>(),
        (int)batch, (int)dim, (float)margin, (float)eps, stream);
    tml_mean_kernel<<<1, TML_BLOCK, 0, stream>>>(
        losses.data_ptr<float>(), out.data_ptr<float>(), (int)batch);

    return {out, dist_pos, dist_neg};
}

std::vector<torch::Tensor> triplet_backward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    torch::Tensor dist_pos,
    torch::Tensor dist_neg,
    torch::Tensor grad_out,
    double margin,
    double eps)
{
    auto a = anchor.contiguous();
    auto p = positive.contiguous();
    auto n = negative.contiguous();
    auto go = grad_out.contiguous();

    const int64_t batch = a.size(0);
    const int64_t dim = batch > 0 ? (a.numel() / batch) : 0;

    auto ga = at::empty_like(a);
    auto gp = at::empty_like(p);
    auto gn = at::empty_like(n);

    if (batch == 0 || dim == 0) {
        return {ga, gp, gn};
    }

    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
    tml_backward_kernel<<<(int)batch, TML_BLOCK, 0, stream>>>(
        a.data_ptr<float>(), p.data_ptr<float>(), n.data_ptr<float>(),
        dist_pos.data_ptr<float>(), dist_neg.data_ptr<float>(), go.data_ptr<float>(),
        ga.data_ptr<float>(), gp.data_ptr<float>(), gn.data_ptr<float>(),
        (int)batch, (int)dim, (float)margin, (float)eps);

    return {ga, gp, gn};
}
"""

_triplet_cpp_src = r"""
#include <torch/extension.h>
#include <vector>

torch::Tensor triplet_forward_eval(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps);

std::vector<torch::Tensor> triplet_forward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    double margin,
    double eps);

std::vector<torch::Tensor> triplet_backward(
    torch::Tensor anchor,
    torch::Tensor positive,
    torch::Tensor negative,
    torch::Tensor dist_pos,
    torch::Tensor dist_neg,
    torch::Tensor grad_out,
    double margin,
    double eps);
"""

_triplet_ext = load_inline(
    name="fused_triplet_margin_loss_v3",
    cpp_sources=_triplet_cpp_src,
    cuda_sources=_triplet_cuda_src,
    functions=["triplet_forward_eval", "triplet_forward", "triplet_backward"],
    verbose=False,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3"],
)


class _FusedTripletMarginLoss(torch.autograd.Function):
    @staticmethod
    def forward(ctx, anchor, positive, negative, margin, eps):
        out, dist_pos, dist_neg = _triplet_ext.triplet_forward(
            anchor, positive, negative, margin, eps
        )
        ctx.save_for_backward(anchor, positive, negative, dist_pos, dist_neg)
        ctx.margin = margin
        ctx.eps = eps
        return out

    @staticmethod
    def backward(ctx, grad_output):
        anchor, positive, negative, dist_pos, dist_neg = ctx.saved_tensors
        ga, gp, gn = _triplet_ext.triplet_backward(
            anchor,
            positive,
            negative,
            dist_pos,
            dist_neg,
            grad_output.contiguous(),
            ctx.margin,
            ctx.eps,
        )
        return ga, gp, gn, None, None


class ModelNew(nn.Module):
    """
    Triplet Margin Loss implemented with a fused custom CUDA kernel.

    Parameters:
        margin (float): The margin between the positive and negative samples.
    """
    def __init__(self, margin=1.0):
        super(ModelNew, self).__init__()
        self.margin = float(margin)
        self.eps = 1e-6  # matches torch.nn.TripletMarginLoss default

    def forward(self, anchor, positive, negative):
        if (
            anchor.is_cuda
            and positive.is_cuda
            and negative.is_cuda
            and anchor.dtype == torch.float32
            and positive.dtype == torch.float32
            and negative.dtype == torch.float32
        ):
            a = anchor if anchor.is_contiguous() else anchor.contiguous()
            p = positive if positive.is_contiguous() else positive.contiguous()
            n = negative if negative.is_contiguous() else negative.contiguous()
            if torch.is_grad_enabled() and (
                a.requires_grad or p.requires_grad or n.requires_grad
            ):
                return _FusedTripletMarginLoss.apply(
                    a, p, n, self.margin, self.eps
                )
            # Fast eval path: no autograd bookkeeping, no distance buffers.
            return _triplet_ext.triplet_forward_eval(
                a, p, n, self.margin, self.eps
            )
        # Fallback for non-CUDA / non-fp32 inputs.
        return F.triplet_margin_loss(
            anchor, positive, negative, margin=self.margin, p=2.0, eps=self.eps
        )
```