Checkbox transition stops working if the checked state is managed externally

[elokour]

What package within Headless UI are you using?
@headlessui/react

What version of that package are you using?
v2.0

What browser are you using?
chrome

Reproduction URL
had some issues with codesandbox so could not provide link
even tried stackblitz and also had issues there.

here is a code snippet of my issue:
my code:

import { Checkbox } from "@headlessui/react";

export default function CheckboxElement({
  checked,
  onChange,
}: {
  checked: boolean;
  onChange: (checked: boolean) => void;
}) {
  return (
    <Checkbox
      checked={checked}
      onChange={onChange}
      className="group block size-4 rounded border bg-transparent transition duration-200 data-[checked]:bg-primary"
    >
      {/* Checkmark icon */}
        <svg
          className="stroke-white opacity-0 transition duration-200 group-data-[checked]:opacity-100"
          viewBox="0 0 14 14"
          fill="none"
        >
          <path
            d="M3 8L6 11L11 3.5"
            strokeWidth={2}
            strokeLinecap="round"
            strokeLinejoin="round"
          />
        </svg>
    </Checkbox>
  );
}

documentation example:

import { Checkbox } from '@headlessui/react'
import { useState } from 'react'

function Example() {
  const [enabled, setEnabled] = useState(false)

  return (
    <Checkbox
      checked={enabled}
      onChange={setEnabled}
      className="group block size-4 rounded border bg-white transition data-[checked]:bg-blue-500"
    >
      <svg
        className="stroke-white opacity-0 transition group-data-[checked]:opacity-100"
        viewBox="0 0 14 14"
        fill="none"
      >
        <path d="M3 8L6 11L11 3.5" strokeWidth={2} strokeLinecap="round" strokeLinejoin="round" />
      </svg>
    </Checkbox>
  )
}

Describe your issue
Because my checked and onChange states are coming from the parent component as props the CSS transition won't work and as soon as I write it like in the documentation where the state is internal the CSS transition works as expected.