KNN with tensorflow

data.json

[
  {
    "x": 1,
    "y": 1,
    "label": "A"
  },
  {
    "x": 1,
    "y": 2,
    "label": "A"
  },
  {
    "x": 2,
    "y": 1,
    "label": "A"
  },
  {
    "x": 2,
    "y": 2,
    "label": "A"
  },
  {
    "x": 1,
    "y": 6,
    "label": "B"
  },
  {
    "x": 2,
    "y": 6,
    "label": "B"
  },
  {
    "x": 5,
    "y": 5,
    "label": "C"
  },
  {
    "x": 5,
    "y": 6,
    "label": "C"
  },
  {
    "x": 6,
    "y": 5,
    "label": "C"
  },
  {
    "x": 6,
    "y": 6,
    "label": "C"
  }
]

tensorflow-knn.js

const tf = require('@tensorflow/tfjs-node')
const data = require('./data.json')

const xs = tf.tensor(data.map(d => d.x))
const ys = tf.tensor(data.map(d => d.y))
const labels = data.map(d => d.label)

const xInput = tf.scalar(Math.random() * xs.max().dataSync())
const yInput = tf.scalar(Math.random() * ys.max().dataSync())

const distances = xInput
  .sub(xs)
  .square()
  .add(
    yInput
    .sub(ys)
    .square()
  )
  .sqrt()

const minorDistance = distances.min().dataSync()[0]
const index = distances.dataSync().indexOf(minorDistance)

console.table({
  x: xInput.arraySync(),
  y: yInput.arraySync(),
  label: labels[index]
})