The first time you run this it will download the dataset, after that it will just load it from the file.

path = Path("~/data/tensorflow/translation/").expanduser()

data_set = "opus/medical"
# data_set = "para_crawl/ende"

train_stream_fn = trax.data.TFDS(data_set,
                                 data_dir=path,
                                 keys=('en', 'de'),
                                 eval_holdout_size=0.01,
                                 train=True)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-6-fb62d04026f5> in <module>
      4 # data_set = "para_crawl/ende"
      5 
----> 6 train_stream_fn = trax.data.TFDS(data_set,
      7                                  data_dir=path,
      8                                  keys=('en', 'de'),

/usr/local/lib/python3.8/dist-packages/gin/config.py in gin_wrapper(*args, **kwargs)
   1067       scope_info = " in scope '{}'".format(scope_str) if scope_str else ''
   1068       err_str = err_str.format(name, fn_or_cls, scope_info)
-> 1069       utils.augment_exception_message_and_reraise(e, err_str)
   1070 
   1071   return gin_wrapper

/usr/local/lib/python3.8/dist-packages/gin/utils.py in augment_exception_message_and_reraise(exception, message)
     39   proxy = ExceptionProxy()
     40   ExceptionProxy.__qualname__ = type(exception).__qualname__
---> 41   raise proxy.with_traceback(exception.__traceback__) from None
     42 
     43 

/usr/local/lib/python3.8/dist-packages/gin/config.py in gin_wrapper(*args, **kwargs)
   1044 
   1045     try:
-> 1046       return fn(*new_args, **new_kwargs)
   1047     except Exception as e:  # pylint: disable=broad-except
   1048       err_str = ''

/usr/local/lib/python3.8/dist-packages/gin/config.py in gin_wrapper(*args, **kwargs)
   1067       scope_info = " in scope '{}'".format(scope_str) if scope_str else ''
   1068       err_str = err_str.format(name, fn_or_cls, scope_info)
-> 1069       utils.augment_exception_message_and_reraise(e, err_str)
   1070 
   1071   return gin_wrapper

/usr/local/lib/python3.8/dist-packages/gin/utils.py in augment_exception_message_and_reraise(exception, message)
     39   proxy = ExceptionProxy()
     40   ExceptionProxy.__qualname__ = type(exception).__qualname__
---> 41   raise proxy.with_traceback(exception.__traceback__) from None
     42 
     43 

/usr/local/lib/python3.8/dist-packages/gin/config.py in gin_wrapper(*args, **kwargs)
   1044 
   1045     try:
-> 1046       return fn(*new_args, **new_kwargs)
   1047     except Exception as e:  # pylint: disable=broad-except
   1048       err_str = ''

~/trax/trax/data/tf_inputs.py in TFDS(dataset_name, data_dir, tfds_preprocess_fn, keys, train, shuffle_train, host_id, n_hosts, eval_holdout_size)
    279   else:
    280     subsplit = None
--> 281   (train_data, eval_data, _) = _train_and_eval_dataset(
    282       dataset_name, data_dir, eval_holdout_size,
    283       train_shuffle_files=shuffle_train, subsplit=subsplit)

~/trax/trax/data/tf_inputs.py in _train_and_eval_dataset(dataset_name, data_dir, eval_holdout_size, train_shuffle_files, eval_shuffle_files, subsplit)
    224   if eval_holdout_examples > 0 or subsplit is not None:
    225     n_train = train_examples - eval_holdout_examples
--> 226     train_start = int(n_train * subsplit[0])
    227     train_end = int(n_train * subsplit[1])
    228     if train_end - train_start < 1:

TypeError: 'NoneType' object is not subscriptable
  In call to configurable 'TFDS' (<function TFDS at 0x7f960c527280>)
  In call to configurable 'TFDS' (<function TFDS at 0x7f960c526f70>)

Since we already downloaded the data in the previous code-block, this will just load the evaluation set from the downloaded data.

eval_stream_fn = trax.data.TFDS('opus/medical',
                                data_dir=path,
                                keys=('en', 'de'),
                                eval_holdout_size=0.01,
                                train=False)

Notice that TFDS returns a generator function, not a generator. This is because in Python, you cannot reset generators so you cannot go back to a previously yielded value. During deep learning training, you use Stochastic Gradient Descent and don't actually need to go back – but it is sometimes good to be able to do that, and that's where the functions come in. Let's print a a sample pair from our train and eval data. Notice that the raw output is represented in bytes (denoted by the b' prefix) and these will be converted to strings internally in the next steps.

train_stream = train_stream_fn()
print(colored('train data (en, de) tuple:', 'red'), next(train_stream))
print()

[31mtrain data (en, de) tuple:[0m (b'Tel: +421 2 57 103 777\n', b'Tel: +421 2 57 103 777\n')

eval_stream = eval_stream_fn()
print(colored('eval data (en, de) tuple:', 'red'), next(eval_stream))

[31meval data (en, de) tuple:[0m (b'Lutropin alfa Subcutaneous use.\n', b'Pulver zur Injektion Lutropin alfa Subkutane Anwendung\n')

We will place a limit on the number of tokens per sentence to ensure we won't run out of memory. This is done with the trax.data.FilterByLength() method and you can see its syntax below.

Filter too long sentences to not run out of memory. length_keys=[0, 1] means we filter both English and German sentences, so both must not be longer that 256 tokens for training and 512 tokens for evaluation.

filtered_train_stream = trax.data.FilterByLength(
    max_length=256, length_keys=[0, 1])(tokenized_train_stream)
filtered_eval_stream = trax.data.FilterByLength(
    max_length=512, length_keys=[0, 1])(tokenized_eval_stream)

train_input, train_target = next(filtered_train_stream)
print(colored(f'Single tokenized example input:', 'red' ), train_input)
print(colored(f'Single tokenized example target:', 'red'), train_target)

[31mSingle tokenized example input:[0m [ 2538  2248    30 12114 23184 16889     5     2 20852  6456 20592  5812
  3932    96  5178  3851    30  7891  3550 30650  4729   992     1]
[31mSingle tokenized example target:[0m [ 1872    11  3544    39  7019 17877 30432    23  6845    10 14222    47
  4004    18 21674     5 27467  9513   920   188 10630    18  3550 30650
  4729   992     1]

Buckets are defined in terms of boundaries and batch sizes. Batch_sizes[i] determines the batch size for items with length < boundaries[i]. So below, we'll take a batch of 256 sentences of length < 8, 128 if length is between 8 and 16, and so on – and only 2 if length is over 512. We'll do the bucketing using bucket_by_length.

boundaries = [2**power_of_two for power_of_two in range(3, 10)]
batch_sizes = [2**power_of_two for power_of_two in range(8, 0, -1)]

Create the generators.

train_batch_stream = trax.data.BucketByLength(
    boundaries, batch_sizes,
    length_keys=[0, 1]  # As before: count inputs and targets to length.
)(filtered_train_stream)

eval_batch_stream = trax.data.BucketByLength(
    boundaries, batch_sizes,
    length_keys=[0, 1]
)(filtered_eval_stream)

Add masking for the padding (0s) using add_loss_weights (we're using AddLossWeights but the documentation for that just says "see add_loss_weights"). I can't find any documentation for it, but I think the 0's are what BucketByLength uses for padding.

train_batch_stream = trax.data.AddLossWeights(id_to_mask=0)(train_batch_stream)
eval_batch_stream = trax.data.AddLossWeights(id_to_mask=0)(eval_batch_stream)

We're going to define a trax function (FN) for addition.

def Addition():
    layer_name = "Addition" 

    def func(x, y):
        return x + y

    return layers.Fn(layer_name, func)

Test it out.

add = Addition()
print(type(add))

<class 'trax.layers.base.PureLayer'>

print("name :", add.name)
print("expected inputs :", add.n_in)
print("promised outputs :", add.n_out)

name : Addition
expected inputs : 2
promised outputs : 1

x = numpy.array([3])
y = numpy.array([4])

print(f"{x} + {y} = {add((x, y))}")

[3] + [4] = [7]

def Multiplication():
    layer_name = "Multiplication"

    def func(x, y):
        return x * y

    return layers.Fn(layer_name, func)

Test it out.

mul = Multiplication()

The properties.

print("name :", mul.name)
print("expected inputs :", mul.n_in)
print("promised outputs :", mul.n_out, "\n")

name : Multiplication
expected inputs : 2
promised outputs : 1 

Some Inputs.

x = numpy.array([7])
y = numpy.array([15])
print("x :", x)
print("y :", y)

x : [7]
y : [15]

The Output

z = mul((x, y))
print(f"{x} * {y} = {mul((x, y))}")

[7] * [15] = [105]

serial = layers.Serial(
    Addition(), Multiplication(), Addition()
)

inputs = (numpy.array([3]), numpy.array([4]), numpy.array([15]), numpy.array([3]))

serial.init(shapes.signature(inputs))

print(serial, "\n")
print("name :", serial.name)
print("sublayers :", serial.sublayers)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out, "\n")

Serial_in4[
  Addition_in2
  Multiplication_in2
  Addition_in2
] 

name : Serial
sublayers : [Addition_in2, Multiplication_in2, Addition_in2]
expected inputs : 4
promised outputs : 1 

print(f"{inputs} -> {serial(inputs)}")

(array([3]), array([4]), array([15]), array([3])) -> [108]

The example with the two simple adition and multiplication functions that where coded together with the serial combinator show how stack semantics work in Trax.

Suppose we want to make the simple calculation \((3 + 4) \times 3 + 4\). We can use Select to perform the calculations in the following manner:

1. input 3 4
2. tl.Select([0, 1, 0, 1])
3. add
4. mul
5. add.

The tl.Select requires a list or tuple of 0-based indices to select elements relative to the top of the stack. For our example, the top of the stack is 3 (which is at index 0) then 4 (index 1) and we us Select to copy the top two elements of the stack and then push all four elements back onto the stack which after the command executes will now contain 3 4 3 4. The steps of the calculation for our example are shown in the table below. As in the previous table each column shows the contents of the stack and the outputs after the operations are carried out.

After processing all the inputs the stack contains 25 which is the result of the calculations.

serial = layers.Serial(
    layers.Select([0, 1, 0, 1]),
    Addition(),
    Multiplication(),
    Addition()
)

Now we'll create the input.

x = (numpy.array([3]), numpy.array([4]))
serial.init(shapes.signature(x))

print(serial, "\n")
print("name :", serial.name)
print("sublayers :", serial.sublayers)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out, "\n")

Serial_in2[
  Select[0,1,0,1]_in2_out4
  Addition_in2
  Multiplication_in2
  Addition_in2
] 

name : Serial
sublayers : [Select[0,1,0,1]_in2_out4, Addition_in2, Multiplication_in2, Addition_in2]
expected inputs : 2
promised outputs : 1 

print(f"{x} -> {serial(x)}")

(array([3]), array([4])) -> [25]

Note that since you are passing in indices to Select, you aren't really using it like a stack, even if behind the scenes it's using push and pop.

serial = layers.Serial(
    layers.Select([2, 1, 1, 2]),
    Addition(),
    Multiplication(),
    Addition()
)

x = (numpy.array([3]), numpy.array([4]), numpy.array([5]))
serial.init(shapes.signature(x))

print(f"{x} -> {serial(x)}")

(array([3]), array([4]), array([5])) -> [41]

print((5 + 4) * 4 + 5)

41

Residual networks (that link is to a research paper, this is wikipedia)are frequently used to make deep models easier to train. Trax already has a built in layer for this. The Residual layer computes the element-wise sum of the stack-top input with the output of the layer series. Let's first see how it is used in the code below:

serial = layers.Serial(
    layers.Select([0, 1, 0, 1]),
    layers.Residual(Addition())
)

print(serial, "\n")
print("name :", serial.name)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out)

Serial_in2_out3[
  Select[0,1,0,1]_in2_out4
  Serial_in2[
    Branch_in2_out2[
      None
      Addition_in2
    ]
    Add_in2
  ]
] 

name : Serial
expected inputs : 2
promised outputs : 3

Here, we use the Serial combinator to define our model. The inputs first goes through a Select layer, followed by a Residual layer which passes the Fn: Addition() layer as an argument. What this means is the Residual layer will take the stack top input at that point and add it to the output of the Fn: Addition() layer. You can picture it like the diagram the below, where x1 and x2 are the inputs to the model:

Now, let's try running our model with some sample inputs and see the result:

x1 = numpy.array([3])
x2 = numpy.array([4])

print(f"{x1} + {x2} -> {serial((x1, x2))}")

[3] + [4] -> (array([10]), array([3]), array([4]))

As you can see, the Residual layer remembers the stack top input (i.e. 3) and adds it to the result of the Fn: Addition() layer (i.e. 3 + 4 = 7). The output of Residual(Addition() is then 3 + 7 = 10 and is pushed onto the stack.

On a different note, you'll notice that the Select layer has 4 outputs but the Fn: Addition() layer only pops 2 inputs from the stack. This means the duplicate inputs (i.e. the 2 rightmost arrows of the Select outputs in the figure above) remain in the stack. This is why you still see it in the output of our simple serial network (i.e. array([3]), array([4])). This is useful if you want to use these duplicate inputs in another layer further down the network.

To strengthen your understanding, you can modify the network above and examine the outputs you get. For example, you can pass the Fn: Multiplication() layer instead in the Residual block:

serial = layers.Serial(
    layers.Select([0, 1, 0, 1]), 
    layers.Residual(Multiplication())
)

print(serial, "\n")
print("name :", serial.name)
print("expected inputs :", serial.n_in)
print("promised outputs :", serial.n_out)

Serial_in2_out3[
  Select[0,1,0,1]_in2_out4
  Serial_in2[
    Branch_in2_out2[
      None
      Multiplication_in2
    ]
    Add_in2
  ]
] 

name : Serial
expected inputs : 2
promised outputs : 3

This means you'll have a different output that will be added to the stack top input saved by the Residual block. The diagram becomes like this:

And you'll get 3 + (3 * 4) = 15 as output of the Residual block:

x1 = numpy.array([3])
x2 = numpy.array([4])

y = serial((x1, x2))
print(f"{x1} * {x2} -> {serial((x1, x2))}")

[3] * [4] -> (array([15]), array([3]), array([4]))

We can express the BLEU score as:

\[ BLEU = BP\left(\prod_{i=1}^{4}precision_i\right)^{(1/4)} \]

with the Brevity Penalty and precision defined as:

\[ BP = min\left(1, e^{(1-(\textit{reference}/\textit{candidate}))}\right) \]

\[ precision_i = \frac {\sum_{snt \in{cand}}\sum_{i\in{snt}}min\Bigl(m^{i}_{cand}, m^{i}_{ref}\Bigr)}{w^{i}_{t}} \]

where:

• \(m^{i}_{cand}\), is the count of i-gram in candidate matching the reference translation.
• \(m^{i}_{ref}\), is the count of i-gram in the reference translation.
• \(w^{i}_{t}\), is the total number of i-grams in candidate translation.

And now for a meaningless plot.

data = pandas.DataFrame.from_dict({"1-gram": [0.8],
                                   "2-gram": [0.7],
                                   "3-gram": [0.6],
                                   "4-gram": [0.5]})
plot = data.hvplot.bar(title="N-Gram Precision").opts(
    width=PLOT.width,
    height=PLOT.height,
    fontscale=PLOT.fontscale,
)

output = Embed(plot=plot, file_name="n_gram_precision")()

print(output)

The n-gram precision counts how many unigrams, bigrams, trigrams, and four-grams (i=1,…,4) match their n-gram counterpart in the reference translations. This term acts as a precision metric. Unigrams account for adequacy while longer n-grams account for fluency of the translation. To avoid overcounting, the n-gram counts are clipped to the maximal n-gram count occurring in the reference (\(m_{n}^{ref}\)). Typically precision shows exponential decay with the with the degree of the n-gram.

Another meaningless plot.

data = pandas.DataFrame.from_dict({"1-gram": [0.8],
                                   "2-gram": [0.77],
                                   "3-gram": [0.74],
                                   "4-gram": [0.71]})
plot = data.hvplot.bar(title="Modified N-Gram Precision").opts(
    width=PLOT.width,
    height=PLOT.height,
    fontscale=PLOT.fontscale
)

output = Embed(plot=plot, file_name="modified_n_gram_precision")()

print(output)

When the n-gram precision is multiplied by the BP, then the exponential decay of n-grams is almost fully compensated. The BLEU score corresponds to a geometric average of this modified n-gram precision.

In this example we will have a reference translation and 2 candidates translations. We will tokenize all sentences using the NLTK.

In this section, we will show a simple pipeline for evaluating machine translated text. Due to storage and speed constraints, we will not be using our own model in this lab. Instead, we will be using Google Translate to generate English to German translations and we will evaluate it against a known evaluation set. There are three files we will need:

1. A source text in English. In this lab, we will use the first 1671 words of the wmt19 evaluation dataset downloaded via SacreBLEU. We just grabbed a subset because of limitations in the number of words that can be translated using Google Translate.
2. A reference translation to German of the corresponding first 1671 words from the original English text. This is also provided by SacreBLEU.
3. A candidate machine translation to German from the same 1671 words. This is generated by feeding the source text to a machine translation model. As mentioned above, we will use Google Translate to generate the translations in this file.

With that, we can now compare the reference an candidate translation to get the BLEU Score.

Load the raw data.

with Path(os.environ["WMT19_SOURCE"]).expanduser().open(encoding="utf-8") as reader:
    wmt19_src_1 = reader.read()

with Path(os.environ["WMT19_REFERENCE"]).expanduser().open(encoding="utf-8") as reader:
    wmt19_ref_1 = reader.read()

with Path(os.environ["WMT19_CANDIDATE"]).expanduser().open(encoding="utf-8") as reader:
    wmt19_can_1 = reader.read()

tokenized_corpus_src = nltk.word_tokenize(wmt19_src_1.lower())
tokenized_corpus_ref = nltk.word_tokenize(wmt19_ref_1.lower())
tokenized_corpus_cand = nltk.word_tokenize(wmt19_can_1.lower())    

Inspecting the first sentence of the data.

print("English source text:\n")
print(f"{wmt19_src_1[0:170]} -> {tokenized_corpus_src[0:30]}\n")
print("German reference translation:\n")
print(f"{wmt19_ref_1[0:219]} -> {tokenized_corpus_ref[0:35]}\n")
print("German machine translation:\n")
print(f"{wmt19_can_1[0:199]} -> {tokenized_corpus_cand[0:29]}")

English source text:

Welsh AMs worried about 'looking like muppets'
There is consternation among some AMs at a suggestion their title should change to MWPs (Member of the Welsh Parliament).
 -> ['\ufeffwelsh', 'ams', 'worried', 'about', "'looking", 'like', "muppets'", 'there', 'is', 'consternation', 'among', 'some', 'ams', 'at', 'a', 'suggestion', 'their', 'title', 'should', 'change', 'to', 'mwps', '(', 'member', 'of', 'the', 'welsh', 'parliament', ')', '.']

German reference translation:

Walisische Ageordnete sorgen sich "wie Dödel auszusehen"
Es herrscht Bestürzung unter einigen Mitgliedern der Versammlung über einen Vorschlag, der ihren Titel zu MWPs (Mitglied der walisischen Parlament) ändern soll.
 -> ['\ufeffwalisische', 'ageordnete', 'sorgen', 'sich', '``', 'wie', 'dödel', 'auszusehen', "''", 'es', 'herrscht', 'bestürzung', 'unter', 'einigen', 'mitgliedern', 'der', 'versammlung', 'über', 'einen', 'vorschlag', ',', 'der', 'ihren', 'titel', 'zu', 'mwps', '(', 'mitglied', 'der', 'walisischen', 'parlament', ')', 'ändern', 'soll', '.']

German machine translation:

Walisische AMs machten sich Sorgen, dass sie wie Muppets aussehen könnten
Einige AMs sind bestürzt über den Vorschlag, ihren Titel in MWPs (Mitglied des walisischen Parlaments) zu ändern.
Es ist aufg -> ['walisische', 'ams', 'machten', 'sich', 'sorgen', ',', 'dass', 'sie', 'wie', 'muppets', 'aussehen', 'könnten', 'einige', 'ams', 'sind', 'bestürzt', 'über', 'den', 'vorschlag', ',', 'ihren', 'titel', 'in', 'mwps', '(', 'mitglied', 'des', 'walisischen', 'parlaments']

print(
    "Results reference versus candidate 1 our own BLEU implementation: ",
    round(bleu_score(tokenized_corpus_cand, tokenized_corpus_ref) * 100, 1),
)

Results reference versus candidate 1 our own BLEU implementation:  43.6

print(
    "Results reference versus candidate 1 sacrebleu library BLEU: ",
    round(sacrebleu.corpus_bleu(wmt19_can_1, wmt19_ref_1).score, 1),
)

Results reference versus candidate 1 sacrebleu library BLEU:  43.2

BLEU Score Interpretation on a Corpus

From the table above (taken from here), we can see the translation is high quality (if you see "Hard to get the gist", please open your workspace, delete `wmt19_can.txt` and get the latest version via the Lab Help button). Moreover, the results of our coded BLEU score are almost identical to those of the SacreBLEU package.

batch_size = 256
train_generator = DataGenerator(data.train.question_one, data.train.question_two,
                                batch_size=batch_size)
validation_generator = DataGenerator(data.validate.question_one,
                                     data.validate.question_two,
                                     batch_size=batch_size)
print(f"training question 1 rows: {len(data.train.question_one):,}")
print(f"validation question 1 rows: {len(data.validate.question_one):,}")

training question 1 rows: 89,179
validation question 1 rows: 22,295

Note: I re-ran this next code block so it's actually the second run.

train_steps = 2000
siamese = SiameseModel(len(loader.vocabulary))
training_loop = train_model(siamese.model, triplet_loss_layer, lr_schedule, steps_per_checkpoint=5)

real_stdout = sys.stdout

TIMER.emit = False
TIMER.start()
with TemporaryFile("w") as temp_file:
    sys.stdout = temp_file
    training_loop.run(train_steps)
TIMER.stop()
sys.stdout = real_stdout
print(f"{TIMER.ended - TIMER.started}")

0:19:46.056057

for mode in training_loop.history.modes:
    print(mode)
    print(training_loop.history.metrics_for_mode(mode))

eval
['metrics/TripletLoss']
train
['metrics/TripletLoss', 'training/gradients_l2', 'training/learning_rate', 'training/loss', 'training/steps per second', 'training/weights_l2']

Let's see if the continues going down.

train_steps = 2000
siamese = SiameseModel(len(loader.vocabulary))
training_loop = train_model(siamese.model, triplet_loss_layer, lr_schedule, steps_per_checkpoint=5)

real_stdout = sys.stdout

TIMER.emit = False
TIMER.start()
with TemporaryFile("w") as temp_file:
    sys.stdout = temp_file
    training_loop.run(train_steps)
TIMER.stop()
sys.stdout = real_stdout
print(f"{TIMER.ended - TIMER.started}")

0:17:41.167719

The first step is to compute the cosine similarity matrix or score in the code. This is \(V_1 V_2^T\) which is generated with fastnp.dot.

The clever arrangement of inputs creates the data needed for positive and negative examples without having to run all pair-wise combinations. Because Q1[n] is a duplicate of only Q2[n], other combinations are explicitly created negative examples or Hard Negative examples. The matrix multiplication efficiently produces the cosine similarity of all positive/negative combinations as shown above on the left side of the diagram. 'Positive' are the results of duplicate examples and 'negative' are the results of explicitly created negative examples. The results for our test case are as expected, V1[0]V2[0] match producing '1' while our other 'positive' cases (in green) don't match well, as was arranged. The V2[2] was set to match V1[3] producing a poor match at score[2,2] and an undesired 'negative' case of a '1' shown in grey.

With the similarity matrix (score) we can begin to implement the loss equations. First, we can extract \(\cos(A,P)\) by utilizing fastnp.diagonal. The goal is to grab all the green entries in the diagram above. This is positive in the code.

Next, we will create the closest_negative. This is the nonduplicate entry in V2 that is closest (has largest cosine similarity) to an entry in V1. Each row, n, of score represents all comparisons of the results of Q1[n] vs Q2[x] within a batch. A specific example in our testcase is row score[2,:]. It has the cosine similarity of V1[2] and V2[x]. The closest_negative, as was arranged, is V2[2] which has a score of 1. This is the maximum value of the 'negative' entries (blue entries in the diagram).

To implement this, we need to pick the maximum entry on a row of score, ignoring the 'positive'/green entries. To avoid selecting the 'positive'/green entries, we can make them larger negative numbers. Multiply fastnp.eye(batch_size) with 2.0 and subtract it out of scores. The result is negative_without_positive. Now we can use fastnp.max, row by row (axis=1), to select the maximum which is closest_negative.

Next, we'll create mean_negative. As the name suggests, this is the mean of all the 'negative'/blue values in score on a row by row basis. We can use fastnp.eye(batch_size) and a constant, this time to create a mask with zeros on the diagonal. Element-wise multiply this with score to get just the 'negative values. This is negative_zero_on_duplicate in the code. Compute the mean by using fastnp.sum on negative_zero_on_duplicate for axis=1 and divide it by (batch_size - 1) . This is mean_negative.

Now, we can compute loss using the two equations above and fastnp.maximum. This will form triplet_loss1 and triplet_loss2.

triple_loss is the fastnp.mean of the sum of the two individual losses.

def TripletLossFn(v1: numpy.ndarray, v2: numpy.ndarray,
                  margin: float=0.25) -> jax.interpreters.xla.DeviceArray:
    """Custom Loss function.

    Args:
       v1 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q1.
       v2 (numpy.ndarray): Array with dimension (batch_size, model_dimension) associated to Q2.
       margin (float, optional): Desired margin. Defaults to 0.25.

    Returns:
       jax.interpreters.xla.DeviceArray: Triplet Loss.
    """
    # use fastnp to take the dot product of the two batches (don't forget to transpose the second argument)
    scores = fastnp.dot(v1, v2.T)
    # calculate new batch size
    batch_size = len(scores)
    # use fastnp to grab all postive =diagonal= entries in =scores=
    positive = fastnp.diagonal(scores)  # the positive ones (duplicates)
    # multiply =fastnp.eye(batch_size)= with 2.0 and subtract it out of =scores=
    negative_without_positive = scores - (fastnp.eye(batch_size) * 2.0)
    # take the row by row =max= of =negative_without_positive=. 
    # Hint: negative_without_positive.max(axis = [?])  
    closest_negative = fastnp.max(negative_without_positive, axis=1)
    # subtract =fastnp.eye(batch_size)= out of 1.0 and do element-wise multiplication with =scores=
    negative_zero_on_duplicate = (1.0 - fastnp.eye(batch_size)) * scores
    # use =fastnp.sum= on =negative_zero_on_duplicate= for =axis=1= and divide it by =(batch_size - 1)= 
    mean_negative = fastnp.sum(negative_zero_on_duplicate, axis=1)/(batch_size - 1)
    # compute =fastnp.maximum= among 0.0 and =A=
    # A = subtract =positive= from =margin= and add =closest_negative= 
    triplet_loss1 = fastnp.maximum(0, margin - positive + closest_negative)
    # compute =fastnp.maximum= among 0.0 and =B=
    # B = subtract =positive= from =margin= and add =mean_negative=
    triplet_loss2 = fastnp.maximum(0, (margin - positive) + mean_negative)
    # add the two losses together and take the =fastnp.mean= of it
    triplet_loss = fastnp.mean(triplet_loss1 + triplet_loss2)
    return triplet_loss

v1 = numpy.array([[0.26726124, 0.53452248, 0.80178373],[0.5178918 , 0.57543534, 0.63297887]])
v2 = numpy.array([[ 0.26726124,  0.53452248,  0.80178373],[-0.5178918 , -0.57543534, -0.63297887]])
triplet_loss = TripletLossFn(v2, v1)
print(f"Triplet Loss: {triplet_loss}")

assert triplet_loss == 0.5

Triplet Loss: 0.5

To make a layer out of a function with no trainable variables, use tl.Fn.

from functools import partial
def TripletLoss(margin=0.25):
    triplet_loss_fn = partial(TripletLossFn, margin=margin)
    return layers.Fn('TripletLoss', triplet_loss_fn)